Model doesn't train after pruning

[gkrisp98]

Hi, I am trying to apply the soft pruning algorithm to a face detector using the FPGMPruner from NNI. I am using the following code:

from nni.algorithms.compression.v2.pytorch.pruning import FPGMPruner
config_list = [{
        'sparsity_per_layer' : compress_rates_total[0],
        'op_types' : ['Conv2d'],
    }, {
        'exclude' : True,
        'op_names' : ['loc.0', 'loc.1', 'loc.2', 'loc.3', 'loc.4', 'loc.5',
                    'conf.0', 'conf.1', 'conf.2', 'conf.3', 'conf.4', 'conf.5'
                    ]
    }]

    #prune the model
    pruner = FPGMPruner(net, config_list)
    pruner.compress()
    pruner._unwrap_model()

    # Main loop
    start_time = time.time()

    for epoch in range(args.start_epoch, args.epochs):
        #current_learning_rate = adjust_learning_rate(optimizer, epoch, args.epochs, args.learning_rate, 20, 0.0001) #learning rate scheduler
        current_learning_rate = adjust_learning_rate(optimizer, epoch, args.gammas, args.schedule)
        losses = 0
        train(train_loader, net, criterion, optimizer, epoch, losses, current_learning_rate, compress_rates_total)
        calc(net)

        config_list = [{
            'sparsity_per_layer' : compress_rates_total[epoch], #new config list for each epoch for pruning
            'op_types' : ['Conv2d'],
        }, {
            'exclude' : True,
            'op_names' : ['loc.0', 'loc.1', 'loc.2', 'loc.3', 'loc.4', 'loc.5',
                        'conf.0', 'conf.1', 'conf.2', 'conf.3', 'conf.4', 'conf.5'
                        ]
        }]

        if epoch % args.epoch_prune == 0 or epoch == args.epochs - 1:  #here it prunes
            pruner = FPGMPruner(net, config_list)
            pruner.compress()
            pruner._unwrap_model()
            print('Model pruned')

        torch.save(net.state_dict(), './weights/epoch{}.pth'.format(epoch))

After inspecting every model that's been produced at each epoch I realized that the train() function is not updating none of the weights of the model except those on the layers I excluded from pruning ('loc.0', 'loc.1' , etc). Does anybody know why is this happening and how I can fix it ?