Hello,
I’m trying to speed up my model inference. It’s a PyTorch module, pretty standard - no special ops, just PyTorch convolution layers.
The export code is copied from this tutorial (optional) Exporting a Model from PyTorch to ONNX and Running it using ONNX Runtime — PyTorch Tutorials 1.9.0+cu102 documentation :
if __name__ == '__main__':
model_str_or_path = sys.argv[1]
model_path = Path(model_str_or_path).expanduser()
# when path exists, we assume its a custom model saved locally
if model_path.exists():
with open(Path(model_path, "separator.json"), "r") as stream:
enc_conf = json.load(stream)
xumx_model, model_nsgt, jagged_slicq_sample = load_target_models(
model_str_or_path=model_path, pretrained=True, sample_rate=enc_conf["sample_rate"], device="cpu"
)
xumx_model.eval()
# Input to the model
torch_out = xumx_model(jagged_slicq_sample)
# Export the model
torch.onnx.export(xumx_model, # model being run
jagged_slicq_sample, # model input (or a tuple for multiple inputs)
"xumx_slicq.onnx", # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
opset_version=11, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
input_names = ['input'], # the model's input names
output_names = ['output'], # the model's output names
dynamic_axes={'input' : {0 : 'batch_size'}, # variable length axes
'output' : {0 : 'batch_size'}})
My model has ~6.7 million parameters (torchinfo output):
===============================================================================================
Layer (type:depth-idx) Output Shape Param #
===============================================================================================
OpenUnmix -- --
Total params: 6,669,912
Trainable params: 6,669,912
Non-trainable params: 0
Total mult-adds (G): 194.27
===============================================================================================
Input size (MB): 28.63
Forward/backward pass size (MB): 9359.33
Params size (MB): 26.68
Estimated Total Size (MB): 9414.64
The trained size for the pth weights file is 28M.
My export command has been running for 3 hours, and has not completed yet. I’m wondering if that’s an expected export time for my size of model.
Thanks.
