0

After trained a module by Detectron2, I tried to export the model to TorchScript, Then I got the following errors:

Could not export Python function call '_ScaleGradient'. Remove calls to Python functions > before export. Did you forget to add @script or @script_method annotation? If this is a > nn.ModuleList, add it to __constants__

I found code is in detectron2/modeling/roi_heads/cascade_rcnn.py

class _ScaleGradient(Function):
    @staticmethod
    def forward(ctx, input, scale):
        ctx.scale = scale
        return input

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output * ctx.scale, None

So I change @statcmethod annos to @torch.jit.script_method, after that, I got a "'ScriptMethodStub' object is not callable" error.

I'm not familar with torchscript, how to fix this problem?

Thanks in advance.

Ben
  • 9
  • 4

1 Answers1

0

It seems don't need that _ScaleGradient method when at inference stage, so I just add the following code to cacasde_rcnn.py

if self.training:
    #call _ScaleGradient.apply
else:
    #don't call _ScaleGradient.apply
Ben
  • 9
  • 4