mmtracking
mmtracking copied to clipboard
how to convert to ONNX ?
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.
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) `