How to use the vscode's debug mode to the Vision Transformer model implemented in timm?

Viewed 76

I'm looking at the implementation of the ViT model in timm, hoping to understand the transfer process of the internal module parameters of this model. But I encountered difficulties when using vscode's debugging mode (I have set "justMyCode": false, so I can jump to other packages).

First of all, let me show you the timm code structure mainly used in this question. I observed that the way that timm generates models is similar to torchvision, except that timm first uses Python Decorators to store the model generation function in a dict, and then calls the function in the dict to generate the model:

@register_model   # Decorates
def vit_base_patch16_224(pretrained=False, **kwargs):
    model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)
    model = _create_vision_transformer('vit_base_patch16_224', pretrained=pretrained, **model_kwargs)
    return model    # generate ViT

You can see that only the dict assignment operation is done in the decorator, which corresponds to a string name for each model generation function, and saves it in the dictionary _model_entrypoints:

_model_entrypoints = {}  # mapping of model names to entrypoint fns

def register_model(fn):
    ......
    model_name = fn.__name__
    _model_entrypoints[model_name] = fn
    ......
    return fn

Next are the specific problems I encountered. First: Why does print() in the source code of the ViT model not take effect (I added print() to both __init__ and forward respectively)? And why the debug mode cannot enter the class VisionTransformer(nn.Module): that has a breakpoint?

class VisionTransformer(nn.Module):   # original code in timm/models/vision_transformer.py
    def __init__(self,......)
        super().__init__()
        print("vit: __init__")  # I added
        ...
        ...
    def forward(self, x):
        print("vit: forward")  # I added
        ...

Second: As shown in the following code, I tried to change the statement that calls the dictionary to the statement that calls the function directly in the create_model() of factory.py. There should be no difference in principle, but an error of __init__() got an unexpected keyword argument will be reported:

def create_model(  # original code in timm/models/factory.py
        model_name,
        pretrained=False,
        checkpoint_path='',
        scriptable=None,
        exportable=None,
        no_jit=None,
        **kwargs):
    ...
    create_fn = model_entrypoint(model_name)   # Call the function in the dict in the previously mentioned decorator
    ...
    with set_layer_config(scriptable=scriptable, exportable=exportable, no_jit=no_jit):
        model = create_fn(pretrained=pretrained, **kwargs)    
        # ↑↑↑ original model generation sentence: The sentence will not report an error.
        # ↓↓↓ I added: The sentence will report got an [unexpected keyword argument 'use_mean_pooling'] error
        model = timm.models.vision_transformer.vit_base_patch16_224(pretrained=pretrained, **kwargs)    
    ...
    return model

The following is my test code, calling create_model() in factory.py:

...
model = create_model(
        'vit_base_patch16_224',
        pretrained=False,
        num_classes=1,
        drop_rate=0.0,
        drop_path_rate=0.1,
        attn_drop_rate=0.0,
        drop_block_rate=None,
        use_mean_pooling=True,
        init_scale=0.001,
)
...
...

I am so confused about the source code of timm. I would be very grateful if any of the above two questions is answered.

0 Answers
Related