[Question] Can CenterTrack model be compiled using TorchScript?

[ievbu]

I want to load model using TorchScript scripting/tracing. However, I unsuccessfully tried:

• scripted_model = torch.jit.script(model) and got error:

  RuntimeError: 
  module has no attribute 'pre_hm_layer':
  at model/networks/dla.py:311:20
  def forward(self, x, pre_img=None, pre_hm=None):
      y = []
      x = self.base_layer(x)
      if pre_img is not None:
          x = x + self.pre_img_layer(pre_img)
      if pre_hm is not None:
          x = x + self.pre_hm_layer(pre_hm)
                  ~~~~~~~~~~~~~~~~~ <--- HERE
      for i in range(6):
          x = getattr(self, 'level{}'.format(i))(x)
          y.append(x)
  
      return y
  '__torch__.tracker.center_track_tracker.model.networks.dla.DLA.forward' is being compiled since it was called from '__torch__.tracker.center_track_tracker.model.networks.dla.DLASeg.imgpre2feats'
  <...>

• traced_model = torch.jit.trace(self.model, torch.rand(1, 3, 512, 512)) and also got error:

  model/networks/DCNv2/dcn_v2.py:31: TracerWarning: Converting a tensor to a Python integer might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  ctx.deformable_groups)
  <...>
  module._c._create_method_from_trace(method_name, func, example_inputs, var_lookup_fn, _force_outplace)
  RuntimeError: Only tensors or tuples of tensors can be output from traced functions (getOutput at /pytorch/torch/csrc/jit/tracer.cpp:209)
  frame #0: c10::Error::Error(c10::SourceLocation, std::string const&) + 0x33 (0x7f72f3f82813 in venv/lib/python3.6/site-packages/torch/lib/libc10.so)
  frame #1: torch::jit::tracer::TracingState::getOutput(c10::IValue const&) + 0x365 (0x7f72925de2a5 in venv/lib/python3.6/site-packages/torch/lib/libtorch.so)
  frame #2: torch::jit::tracer::exit(std::vector<c10::IValue, std::allocator<c10::IValue> > const&) + 0x3c (0x7f72925de5ac in venv/lib/python3.6/site-packages/torch/lib/libtorch.so)
  frame #3: <unknown function> + 0x5705c2 (0x7f72f49455c2 in venv/lib/python3.6/site-packages/torch/lib/libtorch_python.so)

Is it possible in general to save/load CenterTrack using TorchScript scripting/tracing? Maybe I just lack some knowledge of PyTorch.

[xingyizhou]

Sorry for the delayed reply. CenterTrack uses two frames and a prior heatmap as input, you probably want to try something like traced_model = torch.jit.trace(self.model, torch.rand(1, 3, 512, 512), torch.rand(1, 3, 512, 512), torch.rand(1, 1, 512, 512)).

[nuzrub]

Hey @ievbu and @xingyizhou , I got it to work with jit.trace. The main issue is that the model forward function packages the model output within a list. You got to undo that by changing the BaseModel.py, line 91, to return out[-1]. Then, go to the Detector.py process function and remove the [-1] right after the model call. Finally, mark your jit trace as strict=False Here is the trace call.

    dummy_inputs = torch.rand((1, 3, 512, 512)), torch.rand((1, 3, 512, 512)), torch.rand((1, 1, 512, 512))
    self.model = torch.jit.trace(self.model, dummy_inputs, strict=False)

It is also fairly easy to update the project to the current PyTorch version and compile to half precision AND optimized_for_inference. I got a 10-15% speedup from this on a GTX 2070.