Making modules Jit-/Traceable & exportable to ONNX

Author: BloodAxe

pytorch_toolbelt/modules/encoders/common.py

@@ -102,32 +102,33 @@ def forward(self, x: Tensor) -> List[Tensor]:  # skipcq: PYL-W0221
         return _take(output_features, self._layers)
 
     @property
+    @torch.jit.unused
     def channels(self) -> List[int]:
         return self._output_filters
 
     @property
+    @torch.jit.unused
     def strides(self) -> List[int]:
         return self._output_strides
 
     @property
+    @torch.jit.unused
     @pytorch_toolbelt_deprecated("This property is deprecated, please use .strides instead.")
     def output_strides(self) -> List[int]:
         return self.strides
 
     @property
+    @torch.jit.unused
     @pytorch_toolbelt_deprecated("This property is deprecated, please use .channels instead.")
     def output_filters(self) -> List[int]:
         return self.channels
 
-    @property
-    @pytorch_toolbelt_deprecated("This property is deprecated, please don't use it")
-    def encoder_layers(self) -> List[nn.Module]:
-        raise NotImplementedError
-
+    @torch.jit.unused
     def set_trainable(self, trainable):
         for param in self.parameters():
             param.requires_grad = bool(trainable)
 
+    @torch.jit.unused
     def change_input_channels(self, input_channels: int, mode="auto"):
         """
         Change number of channels expected in the input tensor. By default,

pytorch_toolbelt/modules/encoders/timm/efficient_net.py

@@ -87,6 +87,7 @@ def forward(self, x):
         features = self.encoder(x)
         return _take(features, self._layers)
 
+    @torch.jit.unused
     def change_input_channels(self, input_channels: int, mode="auto", **kwargs):
         self.encoder.conv_stem = make_n_channel_input_conv2d_same(
             self.encoder.conv_stem, input_channels, mode, **kwargs

pytorch_toolbelt/modules/encoders/timm/resnet.py

@@ -1,6 +1,7 @@
 from collections import OrderedDict
 from typing import List
 
+import torch
 from torch import nn
 
 from .common import GenericTimmEncoder
@@ -29,6 +30,7 @@ def __init__(self, pretrained=True, layers=None, activation=ACT_RELU):
         self.layer4 = encoder.body.layer4
 
     @property
+    @torch.jit.unused
     def encoder_layers(self) -> List[nn.Module]:
         return [self.stem, self.layer1, self.layer2, self.layer3, self.layer4]
 

tests/test_onnx.py

@@ -0,0 +1,76 @@
+import pytest
+import torch
+
+import pytorch_toolbelt.modules.encoders as E
+from pytorch_toolbelt.modules.backbone.inceptionv4 import inceptionv4
+from pytorch_toolbelt.utils.torch_utils import maybe_cuda, count_parameters
+from pytorch_toolbelt.modules.encoders import timm
+
+skip_if_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="Cuda is not available")
+
+
+@pytest.mark.parametrize(
+    ["encoder", "encoder_params"],
+    [
+        [timm.B0Encoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.MixNetXLEncoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.SKResNet18Encoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.SWSLResNeXt101Encoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.TimmResnet200D, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.HRNetW18Encoder, {"pretrained": False}],
+        [timm.DPN68Encoder, {"pretrained": False}],
+        [timm.NFNetF0Encoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.NFNetF0SEncoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.NFRegNetB0Encoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.TimmRes2Next50Encoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+    ],
+)
+@skip_if_no_cuda
+def test_onnx_export(encoder, encoder_params):
+    import onnx
+
+    model = encoder(**encoder_params).eval()
+
+    print(model.__class__.__name__, count_parameters(model))
+    print(model.strides)
+    print(model.channels)
+    dummy_input = torch.rand((1, 3, 256, 256))
+    dummy_input = maybe_cuda(dummy_input)
+    model = maybe_cuda(model)
+
+    input_names = ["image"]
+    output_names = [f"feature_map_{i}" for i in range(len(model.channels))]
+
+    torch.onnx.export(model, dummy_input, "tmp.onnx", verbose=True, input_names=input_names, output_names=output_names)
+    model = onnx.load("tmp.onnx")
+    onnx.checker.check_model(model)
+
+
+@pytest.mark.parametrize(
+    ["encoder", "encoder_params"],
+    [
+        [timm.B0Encoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.MixNetXLEncoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.SKResNet18Encoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.SWSLResNeXt101Encoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.TimmResnet200D, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.HRNetW18Encoder, {"pretrained": False}],
+        [timm.DPN68Encoder, {"pretrained": False}],
+        [timm.NFNetF0Encoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.NFNetF0SEncoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.NFRegNetB0Encoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+        [timm.TimmRes2Next50Encoder, {"pretrained": False, "layers": [0, 1, 2, 3, 4]}],
+    ],
+)
+@skip_if_no_cuda
+def test_jit_trace(encoder, encoder_params):
+    model = encoder(**encoder_params).eval()
+
+    print(model.__class__.__name__, count_parameters(model))
+    print(model.strides)
+    print(model.channels)
+    dummy_input = torch.rand((1, 3, 256, 256))
+    dummy_input = maybe_cuda(dummy_input)
+    model = maybe_cuda(model)
+
+    model = torch.jit.trace(model, dummy_input, check_trace=True)