Compatibility Benchmark for Torchvision & Timm

[VainF]

We have established a compatibility benchmark for Torchvision version 0.13.1, with a compatibility rate of 85% (73 out of 85 models supported). Our thanks go to @Serjio42 for the insightful discussion in Issue #119.

Successful Pruning: 73 Models
 ['alexnet', 'vit_b_16', 'vit_b_32', 'vit_l_16', 'vit_l_32', 'vit_h_14', 'convnext_tiny', 'convnext_small', 'convnext_base', 'convnext_large', 'densenet121', 'densenet169', 'densenet201', 'densenet161', 'efficientnet_b0', 'efficientnet_b1', 'efficientnet_b2', 'efficientnet_b3', 'efficientnet_b4', 'efficientnet_b5', 'efficientnet_b6', 'efficientnet_b7', 'efficientnet_v2_s', 'efficientnet_v2_m', 'efficientnet_v2_l', 'googlenet', 'inception_v3', 'mnasnet0_5', 'mnasnet0_75', 'mnasnet1_0', 'mnasnet1_3', 'mobilenet_v2', 'mobilenet_v3_large', 'mobilenet_v3_small', 'regnet_y_400mf', 'regnet_y_800mf', 'regnet_y_1_6gf', 'regnet_y_3_2gf', 'regnet_y_8gf', 'regnet_y_16gf', 'regnet_y_32gf', 'regnet_y_128gf', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d', 'wide_resnet50_2', 'wide_resnet101_2', 'fcn_resnet50', 'fcn_resnet101', 'deeplabv3_resnet50', 'deeplabv3_resnet101', 'deeplabv3_mobilenet_v3_large', 'lraspp_mobilenet_v3_large', 'squeezenet1_0', 'squeezenet1_1', 'vgg11', 'vgg13', 'vgg16', 'vgg19', 'vgg11_bn', 'vgg13_bn', 'vgg16_bn', 'vgg19_bn', 'ssdlite320_mobilenet_v3_large', 'ssd300_vgg16', 'fasterrcnn_resnet50_fpn', 'fasterrcnn_resnet50_fpn_v2', 'fasterrcnn_mobilenet_v3_large_320_fpn', 'fasterrcnn_mobilenet_v3_large_fpn']

Unsuccessful Pruning: 12 Models
 ['fcos_resnet50_fpn', 'raft_large', 'shufflenet_v2_x0_5', 'shufflenet_v2_x1_0', 'shufflenet_v2_x1_5', 'shufflenet_v2_x2_0', 'swin_t', 'swin_s', 'swin_b', 'keypointrcnn_resnet50_fpn', 'maskrcnn_resnet50_fpn_v2', 'retinanet_resnet50_fpn_v2']

In the future, we plan to expand our benchmark by utilizing the popular model zoo for classification, Timm. This page is for future discussion on this topic.

[Serjio42]

Hi. I've tried to make an approximate estimation of number of prunable timm models among all 964 models in its actual version v0.6.13. The key difficulty is that each timm model can require model-specific input, some models require extra installations of different stuff.
That's why I divided all models on 3 groups while checking variants: "prunable", "unprunable", "with problems". The last group can be treated as standard error of the estimation :) For fixing it the more hand work is needed.
So, the results of my evaluation of torch-pruning v1.0.0 and timm v0.6.13 are the following:
prunable=379
unprunable=546
with problems=39
Or when calculate the proportion, compatibility rate = 39.3-43.3%

The evaluation code is the following:

import torch
import timm
import torch_pruning as tp
import copy
import csv
import gc
import re

timm_models = timm.list_models()
example_inputs = torch.randn(1,3,256,256)
imp = tp.importance.MagnitudeImportance(p=2, group_reduction="mean")
prunable_list = []
unprunable_list = []
problem_with_input_shape = []

for i, model_name in enumerate(timm_models):
    if 'rexnet' in model_name or 'sequencer' in model_name:  # pruning process stuck with that architectures - skip them.
        unprunable_list.append(model_name)
        continue
    model = timm.create_model(model_name, pretrained=False).eval()
    try:
        test_output = model(example_inputs)
    except Exception as e:
        # trying to get the input size from the error message
        try:
            exception_text = getattr(e, 'message', str(e))
            size_needed = int(exception_text[-5:-2])
            example_inputs = torch.randn(1,3,size_needed,size_needed)
            test_output = model(example_inputs)
        except Exception as e:
            try:  # trying to get the input size from the model name
                pattern = re.search(".+(\d{3})", model_name)
                if pattern:
                    size = int(pattern.group(1))
                    example_inputs = torch.randn(1,3,size,size)
                test_output = model(example_inputs)
            except Exception as e:
                print(f'i={i}, {model_name}. Wrong input error: {e}')
                problem_with_input_shape.append(model_name)
                continue
    prunable = True
    try:
        pruner = None
        pruner = tp.pruner.MagnitudePruner(
                    model, 
                    example_inputs, 
                    global_pruning=False, # If False, a uniform sparsity will be assigned to different layers.
                    importance=imp, # importance criterion for parameter selection
                    iterative_steps=1, # the number of iterations to achieve target sparsity
                    ch_sparsity=0.05,
                    ignored_layers=[],
                    )
        pruner.step()
        test_output = model(example_inputs)
    except Exception as e:
        prunable = False
    
    if prunable:
        prunable_list.append(model_name)
    else:
        unprunable_list.append(model_name)
    
    del model
    if pruner:
        del pruner
    gc.collect()

[VainF]

Hello @Serjio42, I wanted to express my great appreciation for the work you've done.

It's indeed not easy to prepare inputs for different cases individually. I will try your script and update the benchmark soon. It seems like there is still some work to be done to ensure compatibility. Thank you again for your work!

[VainF]

Hi @Serjio42, I find a solution for the varied input sizes.

In timm, we can fetch the input size using model.default_cfg['input_size']. However, there are so many non-standard layers in timm that are highly customized and thus can not be covered by Torch-Pruning. For example, a customized BatchNorm2dAct is used almost everywhere!!

import torch
import timm
import torch_pruning as tp
import gc
import re


timm_models = timm.list_models()
example_inputs = torch.randn(1,3,224,224)
imp = tp.importance.MagnitudeImportance(p=2, group_reduction="mean")
prunable_list = []
unprunable_list = []
problem_with_input_shape = []

for i, model_name in enumerate(timm_models):
    print("Pruning %s..."%model_name)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    if 'rexnet' in model_name or 'sequencer' in model_name or 'botnet' in model_name:  # pruning process stuck with that architectures - skip them.
        unprunable_list.append(model_name)
        continue
    try:
        model = timm.create_model(model_name, pretrained=False, no_jit=True).eval().to(device)
    except: # out of memory error
        model = timm.create_model(model_name, pretrained=False, no_jit=True).eval()
        device = 'cpu'

    input_size = model.default_cfg['input_size']
    example_inputs = torch.randn(1, *input_size).to(device)
    test_output = model(example_inputs)
    #print(model)
    prunable = True
    try:
        pruner = None
        pruner = tp.pruner.MagnitudePruner(
                        model, 
                        example_inputs, 
                        global_pruning=False, # If False, a uniform sparsity will be assigned to different layers.
                        importance=imp, # importance criterion for parameter selection
                        iterative_steps=1, # the number of iterations to achieve target sparsity
                        ch_sparsity=0.05,
                        ignored_layers=[],
                    )
        pruner.step()
        test_output = model(example_inputs)
    except Exception as e:
        prunable = False
    
    if prunable:
        prunable_list.append(model_name)
    else:
        unprunable_list.append(model_name)
    
    
    print("Prunable: %d models, \n %s\n"%(len(prunable_list), prunable_list))
    print("Unprunable: %d models, \n %s\n"%(len(unprunable_list), unprunable_list))
    #print("Problem with input shape: %d models, \n %s\n"%(len(problem_with_input_shape), problem_with_input_shape))

    #del model
    #if pruner:
    #    del pruner
    #gc.collect()

[Serjio42]

@VainF Thanks for info. The code looks good.
The only thing that is not understandable for me is using the option no_jit , timm documentation says about it:

no_jit (bool) — set layer config so that model doesn’t utilize jit scripted layers

Why is it necessary to enable this option?

About timm, yeah, it's quite predictable that many of SOTA/experimental models would have custom layers (50+ as I see in /timm/
layers) and, as I understand, timm extend basic torch/torchvision classes to match these new architectures. It's likely that majority of them will be obliviated but some of them tend to become wide-used.
The question is to wait when it will be implemented in major torch/torchvision library or start to maintain some most prospective ones?..

As I understand now (without deep dive into pruning code) there is one pruning function per one Op class (such as Conv, Linear, etc.) input and output? And it's needed to implement separate functionality for each new class such as BatchNorm2dAct...

[VainF]

I think BatchNorm2dAct might be OK because there is no additional parameters in this customized nn.BatchNorm2d.