Skip to content

covert pytorch to onnx #312

@newbecoder

Description

@newbecoder

when i try to convert dinov3-vitb16 model to onnx, i come across a problem, the exported onnx model's input have two param: input and t_mask, i wanna know what is the t_masks, and if i wanna deploy the model by tritonserver, how can i write the config.pbtxt? here is my convert code :
`# we take DINOv3 ViT-B
MODEL_NAME = "dinov3_vitb16"

def convert(weights):
device = torch.device("cpu")
print(f"Using device: {device}")

# 1. load model 
original_model = torch.hub.load(
    repo_or_dir=DINOV3_LOCATION,
    model=MODEL_NAME,
    source="local",
    weights=weights,
).to(device)
original_model.eval()

# 2. input
dynamic_batch_size = 1
dummy_input = torch.randn(dynamic_batch_size, 3, 224, 224).to(device)

onnx_model_path = "dinov3_vitb16.onnx"
torch.onnx.export(
    original_model,
    dummy_input,
    onnx_model_path,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        'input': {0: 'batch_size'}, 
        'output': {0: 'batch_size'}
    },
    opset_version=14,
    do_constant_folding=True,
    verbose=False,
    # dynamo=True
)

# 3. 
onnx_model = onnx.load(onnx_model_path)
onnx.checker.check_model(onnx_model)

print("\n input layer info:")
for i, input in enumerate(onnx_model.graph.input):
    dims = [dim.dim_value if dim.dim_value > 0 else -1 
            for dim in input.type.tensor_type.shape.dim]
    print(f" {i}: name='{input.name}', shape={dims}")

print("\n output layer info:")
for i, output in enumerate(onnx_model.graph.output):
    dims = [dim.dim_value if dim.dim_value > 0 else -1 
            for dim in output.type.tensor_type.shape.dim]
    print(f" {i}: name='{output.name}', shape={dims}")
return onnx_model_path

`

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions