Skip to content

dnn graph simplifier: support optional constant inputs when match #24609

@fengyuentau

Description

@fengyuentau

Describe the feature and motivation

In an attention subgraph, there are three Slice operators. Slice operator can have up to 5 inputs: data, starts and ends are required, while axes and steps are optional. In the current dnn graph simplifier, the number of inputs must match what we hardcoded in the code. If we want to support all Slice patterns, we need to make several duplicates whose only difference is the number of inputs of Slice operator.

Example:

Current workaround is a manual modification on VitTrack which removes the steps input.

Additional context

Code to export VIT_B_16:

import torch # version: 1.13.1
import torchvision # version: 0.14.1
from onnxsim import simplify # version: 0.4.33

dummy_input = torch.randn(1, 3, 224, 224, device="cpu")
input_names = [ "input" ]
output_names = [ "output" ]

opset = 17

# get and export model
model = torchvision.models.vit_b_16(weights=torchvision.models.ViT_B_16_Weights.DEFAULT)
torch.onnx.export(model, dummy_input, "vit_b_16.opset{}.onnx".format(opset), verbose=False, export_params=True,
                  opset_version=opset, input_names=input_names, output_names=output_names)
# simplify model
model_simplified, check = simplify("vit_b_16.opset{}.onnx".format(opset))
onnx.save(model_simplified, "vit_b_16.opset{}.sim.onnx".format(opset))

Metadata

Metadata

Labels

Type

No type

Projects

No projects

Milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions