AttributeError: 'PrunerModuleWrapper' object has no attribute 'weight'

[damon-93]

NNI 学生项目问题概述 / General Question of Student Program

请简要概述您的问题 / 观点 :
Short summary about the question / idea :

• 我跑通了nni/examples/model_compress/pruning/model_prune_torch.py这个例子, 使用LevelPruner剪枝模型,但是我想将剪枝的模型保存下来,所以我在main函数中添加了保存模型的代码:

for epoch in range(args.prune_epochs):
        pruner.update_epoch(epoch)
        print('# Epoch {} #'.format(epoch))
        train(model, device, train_loader, optimizer_finetune)
        top1 = test(model, device, test_loader)
        if top1 > best_top1:
            best_top1 = top1
            # Export the best model, 'model_path' stores state_dict of the pruned model,
            # mask_path stores mask_dict of the pruned model
            pruner.export_model(model_path=model_path, mask_path=mask_path)
            #------------------------以下代码是我添加的-------------------------
            apply_compression_results(model, mask_path)
            dummy_input = dummy_input = torch.randn([args.batch_size, 1, 28, 28]).to(device)
            m_speedup = ModelSpeedup(model, dummy_input, mask_path, device)
            m_speedup.speedup_model()
            torch.save(model.state_dict(), "use_mask_small_model.pth")

但是, 运行报错如下:

[2021-05-20 17:07:41] INFO (nni.compression.pytorch.compressor/MainThread) Model state_dict saved to ./checkpoints/pruned_naive_mnist_level.pth
[2021-05-20 17:07:41] INFO (nni.compression.pytorch.compressor/MainThread) Mask dict saved to ./checkpoints/mask_naive_mnist_level.pth
Traceback (most recent call last):
  File "/home/sunil/damon/nni/examples/model_compress/pruning/model_prune_torch.py", line 300, in <module>
    main(args)
  File "/home/sunil/damon/nni/examples/model_compress/pruning/model_prune_torch.py", line 275, in main
    apply_compression_results(model, mask_path)
  File "/home/sunil/damon/nni/nni/compression/pytorch/pruning/apply_compression.py", line 27, in apply_compression_results
    module.weight.data = module.weight.data.mul_(masks[name]['weight'])
  File "/home/sunil/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 576, in __getattr__
    type(self).__name__, name))
AttributeError: 'PrunerModuleWrapper' object has no attribute 'weight'

我想这个应该是我添加的代码存在问题,不是环境的问题,请您看看,感谢,十分感谢

我也尝试过将apply_compression_results(model, mask_path)这句去除,程序能运行完成,但是产生的use_mask_small_model.pth模型大小为3.3M, 而./checkpoints/pretrain_naive_mnist_level.pth大小为1.7M,模型反而变大了? 这是我理解上存在偏差吗?剪枝后的模型应该变小才对.

[J-shang]

@damon-93 Hello, try this may work, but I think will also meet other error:

            model_sp = LeNet().to(device)
            model_sp.load_state_dict(torch.load(model_path))
            apply_compression_results(model_sp, mask_path)
            dummy_input = torch.randn([args.batch_size, 1, 28, 28]).to(device)
            m_speedup = ModelSpeedup(model_sp, dummy_input, mask_path, device)
            m_speedup.speedup_model()
            torch.save(model_sp.state_dict(), "use_mask_small_model.pth")

And seems your code is the old version, you could update nni version and nni example code.

The model uses PrunerModuleWrapper to simulate the pruning effectiveness, so the model will bigger the origin. And after speedup, the model will be real pruned. Because in your code the model is a wrapped model, so it is bigger. Note that speedup needs an origin model structure.

[sergiogcharles]

Hi there, I'm getting a similar error using the LotteryTicketPruner.

@damon-93 Hello, try this may work, but I think will also meet other error:

            model_sp = LeNet().to(device)
            model_sp.load_state_dict(torch.load(model_path))
            dummy_input = torch.randn([args.batch_size, 1, 28, 28]).to(device)
            m_speedup = ModelSpeedup(model_sp, dummy_input, mask_path, device)
            m_speedup.speedup_model()
            torch.save(model_sp.state_dict(), "use_mask_small_model.pth")

Don't you still have to use apply_compression_results to multiply the weights by the corresponding mask?

[gsychi]

Seconded here. I ran basic_pruners_torch.py (https://github.com/microsoft/nni/blob/af929fdb848092a7b225a498e6d82d82cf6babfa/examples/model_compress/pruning/basic_pruners_torch.py) on the other pruners and added compatibility for the LotteryTicketPruner, but reached the same issue as stated above.

I've also changed lottery_torch_mnist_fc.py by adding the ModelSpeedup, and reporting FLOPS and parameters; there is no change when using the suggested solution by @J-shang. (as evident in the screenshot). In contrast, running basic_pruners_torch.py on the default pruners will result in a lesser FLOPS and parameters as desired.

Is there another hack to get through the shortcomings of apply_compression_results/ is there an update coming soon?

[J-shang]

@damon-93 @sergiogcharles @FTdiscovery
Sorry for I lose apply_compression_results() in the above code. And I corrected the code.

            model_sp = LeNet().to(device)
            model_sp.load_state_dict(torch.load(model_path))
            apply_compression_results(model_sp, mask_path)
            dummy_input = torch.randn([args.batch_size, 1, 28, 28]).to(device)
            m_speedup = ModelSpeedup(model_sp, dummy_input, mask_path, device)
            m_speedup.speedup_model()
            torch.save(model_sp.state_dict(), "use_mask_small_model.pth")

Yes, you are right, apply_compression_results() is required.