Encounter type error when running pruning examples

[RizhaoCai]

Traceback

Traceback (most recent call last):
File "main_torch_pruner.py", line 97, in
main()
File "main_torch_pruner.py", line 90, in main
train(model, device, train_loader, optimizer)
File "main_torch_pruner.py", line 31, in train
output = model(data)
File "/home/rizhao/anaconda3/envs/distiller/lib/python3.6/site-packages/torch/nn/modules/module.py", line 541, in call
result = self.forward(*input, **kwargs)
File "main_torch_pruner.py", line 16, in forward
x = F.relu(self.conv1(x))
File "/home/rizhao/anaconda3/envs/distiller/lib/python3.6/site-packages/torch/nn/modules/module.py", line 541, in call
result = self.forward(*input, **kwargs)
File "/home/rizhao/anaconda3/envs/distiller/lib/python3.6/site-packages/nni/compression/torch/compressor.py", line 197, in new_forward
mask = self.calc_mask(layer, config)
File "/home/rizhao/anaconda3/envs/distiller/lib/python3.6/site-packages/nni/compression/torch/pruners.py", line 116, in calc_mask
w_abs = weight.abs() * mask
TypeError: mul(): argument 'other' (position 1) must be Tensor, not dict

Hey NNI developers,

I would like to utilize NNI for pruning. However, when I just installed NNI and ran the demo example from here: https://github.com/microsoft/nni/blob/master/examples/model_compress/main_torch_pruner.py

I came with this error.

How to reproduce it:
python main_torch_pruner.py

nni Environment:

• nni version: 1.3
• nni mode(local|pai|remote): local
• OS: Ubuntu 16.04
• python version: 3.6
• is conda or virtualenv used?: Yes
• is running in docker?: Yes. Nvidia-docker
• PyTorch: 1.3
• CUDA: 10.1

Any idea to solve it? If there is something wrong with the demo, there should be bugs in the code.

[RizhaoCai]

I tracked back to the source code and set a tracepoint at nni/compression/torch/pruners.py", line 116.

Then, I found that mask is a dict, which only has a key 'weight' and mask['weight'] is the mask_dict, of which the type a torch.Tensor and the size is the same as the weight parameter.

After that, I modified the code at nni/compression/torch/pruners.py", line 116
from
w_abs = weight.abs() * mask
to
w_abs = weight.abs() * mask['weight'].

And the program runs successfully.
However, I am not sure if this solution would result in other side effects since I am not familiar with NNI.

[scarlett2018]

@Cjkkkk - if @RizhaoCai 's fix works, shall we encourage him/her submit a PR directly =)?

[Cjkkkk]

Thanks for bringing up this issue. Your fix makes sense and you can submit a PR directly.

[scarlett2018]

I tracked back to the source code and set a tracepoint at nni/compression/torch/pruners.py", line 116.

Then, I found that mask is a dict, which only has a key 'weight' and mask['weight'] is the mask_dict, of which the type a torch.Tensor and the size is the same as the weight parameter.

After that, I modified the code at nni/compression/torch/pruners.py", line 116
from
w_abs = weight.abs() * mask
to
w_abs = weight.abs() * mask['weight'].

And the program runs successfully.
However, I am not sure if this solution would result in other side effects since I am not familiar with NNI.

@RizhaoCai - could you please submit a PR for this fix? thanks.

[RizhaoCai]

I tracked back to the source code and set a tracepoint at nni/compression/torch/pruners.py", line 116.
Then, I found that mask is a dict, which only has a key 'weight' and mask['weight'] is the mask_dict, of which the type a torch.Tensor and the size is the same as the weight parameter.

After that, I modified the code at nni/compression/torch/pruners.py", line 116
from
w_abs = weight.abs() * mask
to
w_abs = weight.abs() * mask['weight'].
And the program runs successfully.
However, I am not sure if this solution would result in other side effects since I am not familiar with NNI.

@RizhaoCai - could you please submit a PR for this fix? thanks.

OK. I will check if there will be any side effects and then submit a PR.