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 forforward.
—> 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?