How to save scripted models?

[maxfreu]

Hi, I'm trying to script and save the models, but get different errors depending on arch and encoder. Example:

torch.jit.script(smp.Unet())

results in:

Click me

Traceback (most recent call last):
  File "/home/max/programs/anaconda3/envs/dr/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3427, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-4-a16aa1320739>", line 1, in <module>
    scripted_model = torch.jit.script(smp.Unet())
  File "/home/max/programs/anaconda3/envs/dr/lib/python3.7/site-packages/torch/jit/__init__.py", line 1516, in script
    return torch.jit._recursive.create_script_module(obj, torch.jit._recursive.infer_methods_to_compile)
  File "/home/max/programs/anaconda3/envs/dr/lib/python3.7/site-packages/torch/jit/_recursive.py", line 318, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/home/max/programs/anaconda3/envs/dr/lib/python3.7/site-packages/torch/jit/_recursive.py", line 372, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/home/max/programs/anaconda3/envs/dr/lib/python3.7/site-packages/torch/jit/__init__.py", line 1900, in _construct
    init_fn(script_module)
  File "/home/max/programs/anaconda3/envs/dr/lib/python3.7/site-packages/torch/jit/_recursive.py", line 353, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, infer_methods_to_compile)
  File "/home/max/programs/anaconda3/envs/dr/lib/python3.7/site-packages/torch/jit/_recursive.py", line 376, in create_script_module_impl
    create_methods_from_stubs(concrete_type, stubs)
  File "/home/max/programs/anaconda3/envs/dr/lib/python3.7/site-packages/torch/jit/_recursive.py", line 292, in create_methods_from_stubs
    concrete_type._create_methods(defs, rcbs, defaults)
  File "/home/max/programs/anaconda3/envs/dr/lib/python3.7/site-packages/torch/jit/_recursive.py", line 602, in compile_unbound_method
    create_methods_from_stubs(concrete_type, (stub,))
  File "/home/max/programs/anaconda3/envs/dr/lib/python3.7/site-packages/torch/jit/_recursive.py", line 292, in create_methods_from_stubs
    concrete_type._create_methods(defs, rcbs, defaults)
RuntimeError: 
Class Identity does not have an __init__ function defined:
  File "/home/max/programs/anaconda3/envs/dr/lib/python3.7/site-packages/segmentation_models_pytorch/encoders/resnet.py", line 49
    def get_stages(self):
        return [
            nn.Identity(),
            ~~~~~~~~~~~ <--- HERE
            nn.Sequential(self.conv1, self.bn1, self.relu),
            nn.Sequential(self.maxpool, self.layer1),
'ResNetEncoder.get_stages' is being compiled since it was called from 'ResNetEncoder.forward'
  File "/home/max/programs/anaconda3/envs/dr/lib/python3.7/site-packages/segmentation_models_pytorch/encoders/resnet.py", line 58
    def forward(self, x):
        stages = self.get_stages()
        ~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        features = []

smp version 0.1.3
pytorch version 1.6.0

pytorch 1.7 fails like this:

Click me

RuntimeError: 
tuple slice indices must be integer constants:
  File "/home/max/programs/anaconda3/envs/dr/lib/python3.7/site-packages/segmentation_models_pytorch/encoders/_base.py", line 18
    def out_channels(self):
        """Return channels dimensions for each tensor of forward output of encoder"""
        return self._out_channels[: self._depth + 1]
               ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

EDIT: Working with torch.jit.trace(...) works as a replacement to circumvent loading the code before loading the weights - therefore closing.

[claverru]

Same happened to me, got the second error when trying to convert to torchscript. I think the issue isn't solved since you got a workaround and not a real solution hence it should be opened again.

[maxfreu]

Yeah, I've just hit a case where the workaround doesn't work (Unet-efficientnet-b0, saving works, loading doesn't).

[LSnyd]

I am facing the same issue. @maxfreu have you solved this by now?

[qubvel]

Have you check scripting with timm encoders?

[maxfreu]

torch version 1.8.1, smp version 0.2.0

torch.jit.script(smp.Unet("tu-efficientnet_b0"))

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/max/programs/anaconda3/envs/dr2/lib/python3.8/site-packages/torch/jit/_script.py", line 942, in script
    return torch.jit._recursive.create_script_module(
  File "/home/max/programs/anaconda3/envs/dr2/lib/python3.8/site-packages/torch/jit/_recursive.py", line 391, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/home/max/programs/anaconda3/envs/dr2/lib/python3.8/site-packages/torch/jit/_recursive.py", line 448, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/home/max/programs/anaconda3/envs/dr2/lib/python3.8/site-packages/torch/jit/_script.py", line 391, in _construct
    init_fn(script_module)
  File "/home/max/programs/anaconda3/envs/dr2/lib/python3.8/site-packages/torch/jit/_recursive.py", line 428, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/home/max/programs/anaconda3/envs/dr2/lib/python3.8/site-packages/torch/jit/_recursive.py", line 403, in create_script_module_impl
    method_stubs = stubs_fn(nn_module)
  File "/home/max/programs/anaconda3/envs/dr2/lib/python3.8/site-packages/torch/jit/_recursive.py", line 651, in infer_methods_to_compile
    stubs.append(make_stub_from_method(nn_module, method))
  File "/home/max/programs/anaconda3/envs/dr2/lib/python3.8/site-packages/torch/jit/_recursive.py", line 52, in make_stub_from_method
    return make_stub(func, method_name)
  File "/home/max/programs/anaconda3/envs/dr2/lib/python3.8/site-packages/torch/jit/_recursive.py", line 37, in make_stub
    ast = get_jit_def(func, name, self_name="RecursiveScriptModule")
  File "/home/max/programs/anaconda3/envs/dr2/lib/python3.8/site-packages/torch/jit/frontend.py", line 271, in get_jit_def
    return build_def(ctx, fn_def, type_line, def_name, self_name=self_name)
  File "/home/max/programs/anaconda3/envs/dr2/lib/python3.8/site-packages/torch/jit/frontend.py", line 293, in build_def
    param_list = build_param_list(ctx, py_def.args, self_name)
  File "/home/max/programs/anaconda3/envs/dr2/lib/python3.8/site-packages/torch/jit/frontend.py", line 320, in build_param_list
    raise NotSupportedError(ctx_range, _vararg_kwarg_err)
torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
  File "/home/max/programs/anaconda3/envs/dr2/lib/python3.8/site-packages/segmentation_models_pytorch/unet/decoder.py", line 108
    def forward(self, *features):
                      ~~~~~~~~~ <--- HERE
    
        features = features[1:]    # remove first skip with same spatial resolution

I think the error message is quite clear here: "NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults". So forward should maybe take a list or tuple as input instead of *features.

[github-actions]

This issue is stale because it has been open 60 days with no activity. Remove stale label or comment or this will be closed in 7 days.

[IvanHahan]

I have the same issue in torch==1.10.2. I'd appreciate to have the scripting logic fixed

[github-actions]

This issue is stale because it has been open 60 days with no activity. Remove stale label or comment or this will be closed in 7 days.

[github-actions]

This issue was closed because it has been stalled for 7 days with no activity.