torch.jit.trace of pruned model fails when L2Norm is involved

[abuvaneswari]

Hello,

I am attempting to use Torch_Pruning to prune SSD model.

Note that I use this fork: https://github.com/dkurt/ssd.pytorch/tree/opencv_support

Once I pruned away some of conv filters in the vgg layers, I get the following error with torch.jit.trace on the pruned model:

$ python prune_TP_git_issue.py --model ssd300_mAP_77.43_v2.pth
. . . File "prune_TP_git_issue.py", line 62, in model_output = torch.jit.trace(model, torch_image) File "/home/tfs/venv_ssd.pytorch/lib/python3.6/site-packages/torch/jit/init.py", line 882, in trace check_tolerance, _force_outplace, _module_class) File "/home/tfs/venv_ssd.pytorch/lib/python3.6/site-packages/torch/jit/init.py", line 1034, in trace_module module._c._create_method_from_trace(method_name, func, example_inputs, var_lookup_fn, _force_outplace) File "/home/tfs/venv_ssd.pytorch/lib/python3.6/site-packages/torch/nn/modules/module.py", line 530, in call result = self._slow_forward(*input, **kwargs) File "/home/tfs/venv_ssd.pytorch/lib/python3.6/site-packages/torch/nn/modules/module.py", line 516, in _slow_forward result = self.forward(*input, **kwargs) File "/nas4/tfs/ssd.pytorch/ssd.py", line 89, in forward s = self.L2Norm(x) File "/home/tfs/venv_ssd.pytorch/lib/python3.6/site-packages/torch/nn/modules/module.py", line 530, in call result = self._slow_forward(*input, **kwargs) File "/home/tfs/venv_ssd.pytorch/lib/python3.6/site-packages/torch/nn/modules/module.py", line 516, in _slow_forward result = self.forward(*input, **kwargs) File "/nas4/tfs/ssd.pytorch/layers/modules/l2norm.py", line 23, in forward out = self.weight.view(1, -1, 1, 1) * x RuntimeError: The size of tensor a (512) must match the size of tensor b (256) at non-singleton dimension 1

Attaching the python file. prune_TP_git_issue.ssd.pytorch.py.zip

[VainF]

Hi @abuvaneswari

L2Norm is a customized layer and thus not supported in this repo. I'm working on this to support customized ops and rules.

[VainF]

Please refer to "examples/customized_layer.py". A detailed tutorial will be uploaded soon.

[abuvaneswari]

@VainF , thanks a lot! Your example and the latest commits were extremely helpful. Now the jit trace succeeds!