Jit script issue. How to use it correctly?

I don’t know what is purpose of torch.jit.script because it does not work with most of models.

Here is the code


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")

And I got error message


NotSupportedError Traceback (most recent call last)
/tmp/ipykernel_15354/2867432502.py in
80
81 detect = Detector(config)
—> 82 my_module = torch.jit.script(detect)
83 my_module.save(“my_module.pt”)

~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/_script.py in script(obj, optimize, _frames_up, _rcb, example_inputs)
1285 obj = call_prepare_scriptable_func(obj)
1286 return torch.jit._recursive.create_script_module(
→ 1287 obj, torch.jit._recursive.infer_methods_to_compile
1288 )
1289

~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/_recursive.py in create_script_module(nn_module, stubs_fn, share_types, is_tracing)
475 if not is_tracing:
476 AttributeTypeIsSupportedChecker().check(nn_module)
→ 477 return create_script_module_impl(nn_module, concrete_type, stubs_fn)
478
479 def create_script_module_impl(nn_module, concrete_type, stubs_fn):

~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/_recursive.py in create_script_module_impl(nn_module, concrete_type, stubs_fn)
487 “”"
488 cpp_module = torch._C._create_module_with_type(concrete_type.jit_type)
→ 489 method_stubs = stubs_fn(nn_module)
490 property_stubs = get_property_stubs(nn_module)
491 hook_stubs, pre_hook_stubs = get_hook_stubs(nn_module)

~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/_recursive.py in infer_methods_to_compile(nn_module)
756 stubs = []
757 for method in uniqued_methods:
→ 758 stubs.append(make_stub_from_method(nn_module, method))
759 return overload_stubs + stubs
760

~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/_recursive.py in make_stub_from_method(nn_module, method_name)
68 # In this case, the actual function object will have the name _forward,
69 # even though we requested a stub for forward.
—> 70 return make_stub(func, method_name)
71
72

~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/_recursive.py in make_stub(func, name)
53 def make_stub(func, name):
54 rcb = _jit_internal.createResolutionCallbackFromClosure(func)
—> 55 ast = get_jit_def(func, name, self_name=“RecursiveScriptModule”)
56 return ScriptMethodStub(rcb, ast, func)
57

~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/frontend.py in get_jit_def(fn, def_name, self_name, is_classmethod)
291 pdt_arg_types = type_trace_db.get_args_types(qualname)
292
→ 293 return build_def(parsed_def.ctx, fn_def, type_line, def_name, self_name=self_name, pdt_arg_types=pdt_arg_types)
294
295 # TODO: more robust handling of recognizing ignore context manager

~/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)
329 py_def.col_offset + len(“def”))
330
→ 331 param_list = build_param_list(ctx, py_def.args, self_name, pdt_arg_types)
332 return_type = None
333 if getattr(py_def, ‘returns’, None) is not None:

~/anaconda3/envs/torch/lib/python3.7/site-packages/torch/jit/frontend.py in build_param_list(ctx, py_args, self_name, pdt_arg_types)
357 expr = py_args.vararg
358 ctx_range = ctx.make_range(expr.lineno, expr.col_offset - 1, expr.col_offset + len(expr.arg))
→ 359 raise NotSupportedError(ctx_range, _vararg_kwarg_err)
360 if len(py_args.kw_defaults) > 0:
361 # kw_defaults is a list of the values for the kwargs (which default to None),

NotSupportedError: Compiled functions can’t take variable number of arguments or use keyword-only arguments with defaults:
File “/tmp/ipykernel_15354/2867432502.py”, line 57
def forward(self, *input):
~~~~~~ <— HERE
score = None
geometry = None

How can I solve this?

It looks like the problem is with the easydict. You might want to try using a regular dict instead.