A code snippet to train a model with MNIST dataset and compress it using pruning with PyTorch. Completed full code can be found on here https://github.com/sahandilshan/Simple-NN-Compression

Simple-MNIST-Model-Compression.py

from torch.nn.utils import prune
pruning_percenatage = 0.40
parameters_to_prune = (
    (pruned_model.fc1, 'weight'),
    (pruned_model.fc2, 'weight'),
    (pruned_model.fc3, 'weight'),
    (pruned_model.fc4, 'weight'),
)

prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=pruning_percenatage, # Specifying the percentage
)


prune.remove(pruned_model.fc1, 'weight')
prune.remove(pruned_model.fc2, 'weight')
prune.remove(pruned_model.fc3, 'weight')
prune.remove(pruned_model.fc4, 'weight')