Jit script issue. How to use it correctly?

Soonmyun_Jang: <p>I don’t know what is purpose of torch.jit.script because it does not work with most of models.</p>
<p>Here is the code</p>
<pre><code class="lang-auto">
import torch
import torch.nn as nn
import easydict
import json
from torch import Tensor
from typing import Dict,Tuple,Optional,List
import torch.nn.functional as F

import logging
import torch.nn as nn
import numpy as np

class BaseModel(nn.Module):
    """
    Base class for all models
    """
    def __init__(self, config):
        super(BaseModel, self).__init__()
        self.config = config
        self.logger = logging.getLogger(self.__class__.__name__)
        print('BASEMODEL')

    def forward(self, *input: torch.Tensor):
        """
        Forward pass logic

        :return: Model output
        """
        raise NotImplementedError

    def summary(self):
        """
        Model summary
        """
        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
        params = sum([np.prod(p.size()) for p in model_parameters])
        self.logger.info('Trainable parameters: {}'.format(params))
        self.logger.info(self)
        
class Detector(BaseModel):

    def __init__(self, config):
        super().__init__(config)
        self.scoreMap = nn.Conv2d(32, 1, kernel_size=1)
        self.geoMap = nn.Conv2d(32, 4, kernel_size=1)
        self.angleMap = nn.Conv2d(32, 1, kernel_size=1)
        self.size = config.size
        self.sigmoid = nn.Sigmoid()
        
        my_module = torch.jit.script(self.scoreMap)
        my_module = torch.jit.script(self.geoMap)
        my_module = torch.jit.script(self.angleMap)
        my_module = torch.jit.script(self.sigmoid)

    def forward(self, *input):
        final,  = input
        
        score = self.scoreMap(final)
        score = self.sigmoid(score)

        geoMap = self.geoMap(final)
        geoMap = self.sigmoid(geoMap) * self.size  # TODO: 640 is the image size

        angleMap = self.angleMap(final)
        angleMap = (self.sigmoid(angleMap) - 0.5) * math.pi / 2 # -pi/2  pi/2

        geometry = torch.cat([geoMap, angleMap], dim=1)


        return score, geometry
    
    
x =  '{ "size":"20"}'
y = json.loads(x)
config = easydict.EasyDict(y)

detect = Detector(config)
my_module = torch.jit.script(detect)
my_module.save("my_module.pt")
</code></pre>
<p>And I got error message</p>
<blockquote>
<hr>
<p>NotSupportedError                         Traceback (most recent call last)<br>
/tmp/ipykernel_15354/2867432502.py in <br>
80<br>
81 detect = Detector(config)<br>
—&gt; 82 my_module = torch.jit.script(detect)<br>
83 my_module.save(“my_module.pt”)</p>
<p>~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/_script.py in script(obj, optimize, _frames_up, _rcb, example_inputs)<br>
1285         obj = call_prepare_scriptable_func(obj)<br>
1286         return torch.jit._recursive.create_script_module(<br>
 → 1287             obj, torch.jit._recursive.infer_methods_to_compile<br>
1288         )<br>
1289</p>
<p>~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/_recursive.py in create_script_module(nn_module, stubs_fn, share_types, is_tracing)<br>
475     if not is_tracing:<br>
476         AttributeTypeIsSupportedChecker().check(nn_module)<br>
 → 477     return create_script_module_impl(nn_module, concrete_type, stubs_fn)<br>
478<br>
479 def create_script_module_impl(nn_module, concrete_type, stubs_fn):</p>
<p>~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/_recursive.py in create_script_module_impl(nn_module, concrete_type, stubs_fn)<br>
487     “”"<br>
488     cpp_module = torch._C._create_module_with_type(concrete_type.jit_type)<br>
 → 489     method_stubs = stubs_fn(nn_module)<br>
490     property_stubs = get_property_stubs(nn_module)<br>
491     hook_stubs, pre_hook_stubs = get_hook_stubs(nn_module)</p>
<p>~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/_recursive.py in infer_methods_to_compile(nn_module)<br>
756     stubs = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
757     for method in uniqued_methods:<br>
 → 758         stubs.append(make_stub_from_method(nn_module, method))<br>
759     return overload_stubs + stubs<br>
760</p>
<p>~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/_recursive.py in make_stub_from_method(nn_module, method_name)<br>
68     # In this case, the actual function object will have the name <code>_forward</code>,<br>
69     # even though we requested a stub for <code>forward</code>.<br>
—&gt; 70     return make_stub(func, method_name)<br>
71<br>
72</p>
<p>~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/_recursive.py in make_stub(func, name)<br>
53 def make_stub(func, name):<br>
54     rcb = _jit_internal.createResolutionCallbackFromClosure(func)<br>
—&gt; 55     ast = get_jit_def(func, name, self_name=“RecursiveScriptModule”)<br>
56     return ScriptMethodStub(rcb, ast, func)<br>
57</p>
<p>~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/frontend.py in get_jit_def(fn, def_name, self_name, is_classmethod)<br>
291         pdt_arg_types = type_trace_db.get_args_types(qualname)<br>
292<br>
 → 293     return build_def(parsed_def.ctx, fn_def, type_line, def_name, self_name=self_name, pdt_arg_types=pdt_arg_types)<br>
294<br>
295 # TODO: more robust handling of recognizing ignore context manager</p>
<p>~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/frontend.py in build_def(ctx, py_def, type_line, def_name, self_name, pdt_arg_types)<br>
329                        py_def.col_offset + len(“def”))<br>
330<br>
 → 331     param_list = build_param_list(ctx, py_def.args, self_name, pdt_arg_types)<br>
332     return_type = None<br>
333     if getattr(py_def, ‘returns’, None) is not None:</p>
<p>~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/frontend.py in build_param_list(ctx, py_args, self_name, pdt_arg_types)<br>
357         expr = py_args.vararg<br>
358         ctx_range = ctx.make_range(expr.lineno, expr.col_offset - 1, expr.col_offset + len(expr.arg))<br>
 → 359         raise NotSupportedError(ctx_range, _vararg_kwarg_err)<br>
360     if len(py_args.kw_defaults) &gt; 0:<br>
361         # kw_defaults is a list of the values for the kwargs (which default to None),</p>
<p>NotSupportedError: Compiled functions can’t take variable number of arguments or use keyword-only arguments with defaults:<br>
File “/tmp/ipykernel_15354/2867432502.py”, line 57<br>
def forward(self, *input):<br>
~~~~~~ &lt;— HERE<br>
score = None<br>
geometry = None</p>
</blockquote>
<p>How can I solve this?</p>

clerkie: <p>It looks like the problem is with the easydict. You might want to try using a regular dict instead.</p>