Sensitivity pruning not setting network weights to zero

[elliottloveridge]

I'm implementing both sensitivity and sparsity level pruning for 3D networks.

Sensitivity analysis works as expected as can be seen from the figure below so I assume it is not a problem with using a 3D model for element-wise pruning methods.

However, when pruning during training the network weights are not going to zero.

My implementation can be abbreviated as the following:

# for each epoch
for i in range(opt.begin_epoch, opt.begin_epoch + opt.n_epochs):

        compression_scheduler.on_epoch_begin(i)

        model.train()

        # for each mini-batch
        for j, (inputs, targets) in enumerate(train_loader):

               compression_scheduler.on_minibatch_begin(i, minibatch_id=j, minibatches_per_epoch=len(train_loader))

                targets = targets.cuda()
                inputs = Variable(inputs)
                targets = Variable(targets)
                outputs = model(inputs)
                loss = criterion(outputs, targets)

                # before backwards pass - update loss to include regularization
                compression_scheduler.before_backward_pass(i, minibatch_id=j, minibatches_per_epoch=len(train_loader), loss=loss)

                optimizer.zero_grad()
                loss.backward()

                compression_scheduler.before_parameter_optimization(i, minibatch_id=j, minibatches_per_epoch=len(train_loader), optimizer=optimizer)

                optimizer.step()

                compression_scheduler.on_minibatch_end(i, minibatch_id=j, minibatches_per_epoch=len(train_loader))

And an example compression file I'm using would be:

version: 1
pruners:
    mobilenetv2_fullyconnected:
        class: 'SensitivityPruner'
        sensitivities:
            'module.features.18.1.weight': 1.0
            'module.classifier.1.weight': 1.0

policies:
    - pruner:
        instance_name: 'mobilenetv2_fullyconnected'
    starting_epoch: 0
    ending_epoch: 20
    frequency: 1

Any help would be appreciated!