hukenovs / hagrid

HAnd Gesture Recognition Image Dataset
https://arxiv.org/abs/2206.08219
589 stars 89 forks source link

Convert pth to onnx #72

Open scottcali opened 5 months ago

scottcali commented 5 months ago

Hi,

Could you write a python script to convert those pth model to onnx?

Thanks!

-Scott

scottcali commented 5 months ago

pytorch_to_onnx.py.txt

I try to write one pytorch_to_onnx.py, attached here, but get the following errors: python3 pytorch_to_onnx.py /home/dev/.local/lib/python3.10/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead. warnings.warn( /home/dev/.local/lib/python3.10/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or None for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing weights=None. warnings.warn(msg) <models.classifiers.base_model.ClassifierModel object at 0x7f60d2d0bb20> exporting model to ONNX... Traceback (most recent call last): File "/mnt/d/Workspace2021/models/hand/hagrid/pytorch_to_onnx.py", line 170, in torch.onnx.export(model, input, "../ResNext50_32.onnx", verbose=True, export_params=True, input_names=input_names, output_names=output_names, opset_version=12) File "/home/dev/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 516, in export _export( File "/home/dev/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 1573, in _export with exporter_context(model, training, verbose): File "/usr/lib/python3.10/contextlib.py", line 135, in enter return next(self.gen) File "/home/dev/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 179, in exporter_context with select_model_mode_for_export( File "/usr/lib/python3.10/contextlib.py", line 135, in enter return next(self.gen) File "/home/dev/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 140, in disable_apex_o2_state_dict_hook for module in model.modules(): AttributeError: 'ClassifierModel' object has no attribute 'modules'

scottcali commented 5 months ago

Try this , same error:

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

input = torch.randn((10, 3, 640, 1280)).float().to(device)

conf = OmegaConf.load("./configs/ResNext50.yaml")
model = build_model(conf)
transform = Demo.get_transform_for_inf(conf.test_transforms)
conf.model.checkpoint = "../ResNext50.pth"
snapshot = torch.load(conf.model.checkpoint, map_location=torch.device("cpu"))
model.load_state_dict(snapshot["MODEL_STATE"])    

#if conf.model.checkpoint is not None:
#    snapshot = torch.load(conf.model.checkpoint, map_location=torch.device("cpu"))
#    model.load_state_dict(snapshot["MODEL_STATE"])    

model.eval()
if model is not None:

    print(model)
    #checkpoint = torch.load("../ResNext50.pth")
    #model.load_state_dict(checkpoint['net'])

    input_names = [ "image" ]
    output_names = [ "output" ]

    print('exporting model to ONNX...')
    torch.onnx.export(model, input, "../ResNext50_32.onnx", verbose=True, export_params=True, input_names=input_names, output_names=output_names, opset_version=12)
wdkwyf commented 3 months ago

@scottcali you should use model.hagrid_model to export.