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.

[hgaiser]

I don't think this issue is resolved, right? It would be good for the models in this repository to be 'scriptable'. The models from timm appear to be scriptable, so there should be no issue there.

@qubvel is this something you would accept help on?

[kargarisaac]

No solution here?

[kargarisaac]

#755 has a solution