Model size after pruning

1

The model is pruned after training. I suppose the pruned model should has less size than original model, right? But it actually become double after pruning. Is there a right way can really reduce model size?

mymodel.pth 72.8MB
my model_pruned.pth 136.1MB

def get_model():
    model = models.vgg16(pretrained=True)
    for param in model.parameters():
        param.requires_grad = False
    model.avgpool = nn.Sequential( nn.Conv2d(512,512,3),
      nn.MaxPool2d(2),
      nn.Flatten())
    model.classifier = nn.Sequential(
      nn.Linear(2048, 512),
      nn.ReLU(),
      nn.Dropout(0.5),
      nn.Linear(512, 136),
      nn.Sigmoid()
    )
    
    criterion = nn.L1Loss()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    return model.to(device), criterion, optimizer

def prun_(model):
    parameters_to_prune = (
       (model.features[0],'weight'),
       (model.features[2],'weight'),
       (model.features[5],'weight'),
       (model.features[7],'weight'),
       (model.features[10],'weight'),
       (model.features[12],'weight'),
       (model.features[14],'weight'),
       (model.features[17],'weight'),
       (model.features[19],'weight'),
       (model.features[21],'weight'),
       (model.features[24],'weight'),
       (model.features[26],'weight'),
       (model.features[28],'weight'),
       (model.classifier[0],'weight'),
       (model.classifier[3],'weight'),
    )
    prune.global_unstructured(
        parameters_to_prune,
        pruning_method=prune.L1Unstructured,
        amount=0.4,
    )

model.eval()
print(model.state_dict().keys())
traced_script_module.save("mymodel.pt")

prun_(model)
print(model.state_dict().keys())
torch.save(model.to('cpu').state_dict(), 'mymodel_pruned.pth')
2

I find the pruned model include additional “original” and “mask” data. It could be the reason that file become larger. Could I throw away these data? Is there a convenient to do that? Thanks!

3

If you want to remove the pruning afterwards you can do torch.nn.utils.prune.remove

4

Finally,I can use [`torch.nn.utils.prune.remove’] to remove the pruning.
The original and pruned model have the same size. I test one forward pass, the original model takes 0.0995msec, and the pruned model takes 0.102msec.
it trun out the model size are the same, but implement time are almost the same. Why pruning? It is quite strange.

1 Like