Torch.onnx.export`s export_params creates separated weight files

jit
1

Hello
When I tried to export torch model to onnx with export_params=True, parameter files for each module in my model saved separately in folder I assigned(./onnx/). It there any way to pack all weights in asr.onnx file?

Thank you

torch.onnx.export(
    model, 
    (torch.randn(3, 80, 355), torch.ones(3, 355)), 
    "./onnx/asr.onnx", 
    export_params=True,
    opset_version=14,
    input_names=["input_features", "feature_mask"],
    output_names=["logits"],
    dynamic_axes={
        "input_features" : {0 : "batch_size", 2 : "sequence_length"},
        "feature_mask" : {0 : "batch_size", 1 : "sequence_length"},
        "logits" : {0 : "batch_size", 1 : "half_sequence_length"}
    }
)

image
image236×726 3.84 KB

2

Same problem with torch.version ‘2.0.1+cu117’

3

Hello, When I inputs torch jit traced model to torch.onnx.export, all weights packed in asr.onnx file.

traced_model = torch.jit.trace(model, [dummy_input_features, dummy_feature_mask])
torch.onnx.export(
    traced_model, 
    (torch.randn(3, 80, 355, dtype=torch.float16, device="cuda:0"), torch.ones(3, 355, dtype=torch.float16, device="cuda:0")), 
    "./onnx/asr.onnx", 
    export_params=True,
    opset_version=14,
    input_names=["input_features", "feature_mask"],
    output_names=["logits"],
    dynamic_axes={
        "input_features" : {0 : "batch_size", 2 : "sequence_length"},
        "feature_mask" : {0 : "batch_size", 1 : "sequence_length"},
        "logits" : {0 : "batch_size", 1 : "half_sequence_length"}
    }
)
4

Jit traced model not works for me. Putting the model to cuda and changing the dtype to float16 (half) works for me as you also did. Maybe that’s the solution.