mmtracking icon indicating copy to clipboard operation
mmtracking copied to clipboard

how to convert to ONNX ?

Open mispa-ms opened this issue 2 years ago • 1 comments

Do you support for conversion to onnx? Do you have a script to do it? What are the steps we need to convert MOT model to onnx? Thank you in advance for your help.

mispa-ms avatar Apr 05 '23 20:04 mispa-ms

May be it help you. `import torch from mmdet.apis import init_detector import mmcv

Step 1: Set up the environment and load the model

config_file = 'vfnet_r50_fpn_1x_coco_coast_guard_1_whole.py' checkpoint_file = 'epoch_20.pth' output_file = 'epoch_20.onnx' input_shape = (1, 3, 1280, 640) # Example input shape

Load the model

model = init_detector(config_file, checkpoint_file)

Step 2: Load the weights and set the model to evaluation mode

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = model.to(device).eval()

Step 3: Export the model to ONNX

input_data = torch.randn(input_shape).to(device) # Create a dummy input tensor onnx_output = output_file.replace('.onnx', '_trace.onnx') # Intermediate traced ONNX file

Export the model to ONNX

torch.onnx.export( model, input_data, onnx_output, input_names=['input'], output_names=['output'], opset_version=11, # ONNX opset version do_constant_folding=True, verbose=True )

Optimize the ONNX model (optional)

optimized_model, _ = mmcv.runner.load_model(onnx_output) optimized_model.forward = None # Remove the forward function to enable ONNX optimization optimized_model.cfg = model.cfg optimized_model.forward_dummy = model.forward_dummy mmcv.runner.save_checkpoint(optimized_model, output_file) `

RoyAn2386 avatar Jun 20 '23 04:06 RoyAn2386