How to load a complete pruned model based on saved MLP weights and Head weights

[Fly-Yun-Maker]

Hello, I would like to know if it is possible to dynamically generate a model with the corresponding compression rate based on the saved weights for training other downstream tasks. May I ask if you can provide the code

[WalterSimoncini]

Hi @Fly-Yun-Maker, sorry for the late reply.

I'm not exactly sure what you meant by your question, did you mean:

1. How to load a pruned model from disk?
2. How to save the importance scores to disk and use them at test-time to produce a pruned model for any arbitrary sparsity?

To answer the first question, you should first load a base model as follows:

import torch

from src.models import load_model, ModelType

cache_dir = "./cache"
device = torch.device("cuda")

model, transform = load_model(
    type_=ModelType.DINO_VIT_B_16,
    cache_dir=cache_dir,
    device=device
)

and then load the pruned checkpoint as follows:

from src.utils.checkpoints import load_pruned_checkpoint

pruned_checkpoint_path = "./path/to/pruned.pt"

model = load_pruned_checkpoint(
    base_model=model,
    checkpoint_path=pruned_checkpoint_path,
    device=device
)

# The model is now pruned, try running the following:
print(sum(p.numel() for p in model.parameters()))

Previously, load_pruned_checkpoint simply sliced the weight matrices, so running print(model) showed the original in_features and out_features values for linear layers, even though the weight matrices were smaller. I've since updated the code in commit 111bd33 to also re-initialize the linear layers with the correct values for those fields.

Regarding the second question, there isn't a clear way to do so at the moment. I will implement some code to do so shortly.