Increace average CPU and GPU memory for pruned model

[bestload]

Hello! Faced with the problem of increasing the amount of CPU and gpu RAM, which is required to pass the test by a convolutional neural network. The code is given below. This may probably be due to attempts to restore the thinned scales. I will be glad to hear exactly the explanation)!!!

Test function for CPU (Similar for gpu use):

import psutil

torch.cuda.empty_cache()

if isinstance(model, nn.DataParallel):
    model = model.module  # Unwrap the DataParallel
    print("DataParallel wrapper removed.")

def test_model_cpu_memory_avg(model, dataloader, device):
    total_memory = 0
    num_batches = 0

    process = psutil.Process()

    model.to(device)
    model.eval()

    with torch.no_grad(), torch.amp.autocast('cpu'):
        for images, labels in dataloader:
            images, labels = images.to(device), labels.to(device)

            start_memory = process.memory_info().rss / 1024 ** 2 

            outputs = model(images)

            end_memory = process.memory_info().rss / 1024 ** 2 

            batch_memory = (start_memory + end_memory) / 2
            total_memory += batch_memory
            num_batches += 1

    avg_memory = total_memory / num_batches

    print(f"Average CPU Memory Usage: {avg_memory:.2f} MB")

test_model_cpu_memory_avg(model, test_loader, 'cpu')

Pruning:

import torch_pruning as tp

if isinstance(model, nn.DataParallel):
    model = model.module
    print("DataParallel wrapper removed.")
    
# Freeze all layers
for param in model.parameters():
    param.requires_grad = True

example_inputs = torch.randn(1, 3, 224, 224)
# Move example_inputs to the same device as the model
example_inputs = example_inputs.to('cpu')  # Move to GPU
importance = tp.importance.GroupNormImportance(p=2)

# model = model.module
model = model.to('cpu')

# ignore final layer
ignored_layers = []
for name, module in model.named_modules():
    if name == 'fc':
        ignored_layers.append(module)

# ignore first layer
for name, module in model.named_modules():
    if name == 'conv1':
        ignored_layers.append(module)      

print(ignored_layers)

pruner = tp.pruner.MetaPruner( # We can always choose MetaPruner if sparse training is not required.
    model,
    example_inputs,
    importance=imp,
    pruning_ratio=0.5, # remove 50% channels, ResNet18 = {64, 128, 256, 512} => ResNet18_Half = {32, 64, 128, 256}
    # pruning_ratio_dict = {model.conv1: 0.2, model.layer2: 0.8}, # customized pruning ratios for layers or blocks
    ignored_layers=ignored_layers,
)

---

And in my case:
Average cpu memory with ResNet-50 for test: 1289.12
Average cpu memory with PRUNED (50%+fine-tuning+50%) ResNet-50 for test: 1536.86 MB