Pytorch DeepLabV3 mobile export to ONNX

Question

I have finetuned different versions of deeplabv3 implemented in pytorch (https://pytorch.org/vision/main/models/deeplabv3.html)

I am looking to export my 3 models to ONNX after testing them on images.

It works very well except on the Mobile version. Yet it is the most important one for me.

Here is the error I get only on the mobile version export:

E:\Documents\Florian\Programmation\QuickTestModel\venv\lib\site-packages\torchvision\io\image.py:11: UserWarning: Failed to load image Python extension: Could not find module 'E:\Documents\Florian\Programmation\QuickTestModel\venv\Lib\site-packages\torchvision\image.pyd' (or one of its dependencies). Try using the full path with constructor syntax.
  warn(f"Failed to load image Python extension: {e}")
Traceback (most recent call last):
  File "E:\Documents\Florian\Programmation\QuickTestModel\export_onnx.py", line 31, in <module>
    torch_out = model(input_batch)
  File "E:\Documents\Florian\Programmation\QuickTestModel\venv\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "E:\Documents\Florian\Programmation\QuickTestModel\venv\lib\site-packages\torchvision\models\segmentation\_utils.py", line 25, in forward
    features = self.backbone(x)
  File "E:\Documents\Florian\Programmation\QuickTestModel\venv\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "E:\Documents\Florian\Programmation\QuickTestModel\venv\lib\site-packages\torchvision\models\_utils.py", line 62, in forward
    x = module(x)
  File "E:\Documents\Florian\Programmation\QuickTestModel\venv\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "E:\Documents\Florian\Programmation\QuickTestModel\venv\lib\site-packages\torchvision\models\mobilenetv3.py", line 89, in forward
    result = self.block(input)
  File "E:\Documents\Florian\Programmation\QuickTestModel\venv\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "E:\Documents\Florian\Programmation\QuickTestModel\venv\lib\site-packages\torch\nn\modules\container.py", line 141, in forward
    input = module(input)
  File "E:\Documents\Florian\Programmation\QuickTestModel\venv\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "E:\Documents\Florian\Programmation\QuickTestModel\venv\lib\site-packages\torchvision\ops\misc.py", line 151, in forward
    scale = self._scale(input)
  File "E:\Documents\Florian\Programmation\QuickTestModel\venv\lib\site-packages\torchvision\ops\misc.py", line 144, in _scale
    scale = self.avgpool(input)
  File "E:\Documents\Florian\Programmation\QuickTestModel\venv\lib\site-packages\torch\nn\modules\module.py", line 1177, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'SqueezeExcitation' object has no attribute 'avgpool'

Process finished with exit code 1

Here the code :

import torch
import torch.onnx

if __name__ == '__main__':

    model = torch.load(r"result/deepllabv3_trace_hum_v2_mobile/weights.pt")
    model = model.to("cpu")

    x = torch.randn(3, 960, 540)
    input_batch = x.unsqueeze(0)


    with torch.no_grad():
        torch_out = model(input_batch)

    torch.onnx.export(model, input_batch, f="model.onnx", export_params=True, opset_version=14,
                      do_constant_folding=True, input_names=['input'],
                  output_names=['output'], dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes
                                'output' : {0 : 'batch_size'}})

Do you have any idea ? Thanks in advance