ALBEF: Train from scratch about multimodal

[ebsmothers]

Hi @XinhaoMei, thanks for using the repo! For ALBEF we only provide the fine-tuning code, so if you want to use the starting checkpoints for pretraining you just need to remap the state dict keys. We often don't keep this code in the repo as it is a bit manual, we instead do it once and then provide our remapped checkpoints (for ALBEF you can find them here, though these ones are already pretrained). Below is an example of how you can do this for pretrained ViT. Lmk if it makes sense, happy to also provide something similar for the BERT mapping.

import torch
from torchmultimodal.models.albef.image_encoder import ALBEFVisionEncoder

# Function to convert pretrained ViT state dict into TorchMultimodal state dict
def map_vit_state_dict(vit_ckpt):
    mapped_state_dict = {}
    encoder_layer_mapping = {
        'norm1': 'ln_1',
        'norm2': 'ln_2',
        'attn.qkv.weight': 'self_attention.in_proj_weight',
        'attn.qkv.bias': 'self_attention.in_proj_bias',
        'attn.proj': 'self_attention.out_proj',
        'norm2': 'ln_2',
        'mlp.fc1': 'mlp.0',
        'mlp.fc2': 'mlp.3'
    }

    misc_mapping = {
        'cls_token': 'encoder.class_token',
        'pos_embed': 'encoder.encoder.pos_embedding',
        'patch_embed.proj': 'encoder.conv_proj',
        'norm': 'encoder.encoder.ln'
    }

    for k, v in vit_ckpt['model'].items():
        if 'blocks' in k:
            layer_num = k.split('.')[1]
            new_k = k.replace('blocks.', 'encoder.encoder.layers.encoder_layer_')
            param_name = [x for x in encoder_layer_mapping.keys() if x in new_k][0]
            new_k = new_k.replace(param_name, encoder_layer_mapping[param_name])
            mapped_state_dict[new_k] = v
        else:
            param_names = [x for x in misc_mapping.keys() if x in k]
            if len(param_names) != 0:
                param_name = param_names[0]
                new_k = k.replace(param_name, misc_mapping[param_name])
                mapped_state_dict[new_k] = v
    return mapped_state_dict

# Load ALBEF vision encoder from TorchMultimodal
tmm_albef_vision_encoder = ALBEFVisionEncoder(
    hidden_size=768,
    image_size=224,
    patch_size=16,
    num_hidden_layers=12,
    num_attention_heads=12,
    mlp_dim=3072,
    dropout=0.0,
    attention_dropout=0.0,
    layer_norm_eps=1e-6
)

# Load pretrained ViT from ALBEF pretraining
vit_ckpt = torch.hub.load_state_dict_from_url(
    url="https://dl.fbaipublicfiles.com/deit/deit_base_patch16_224-b5f2ef4d.pth",
    map_location="cpu", check_hash=True
)

# Map state dict and load into TorchMultimodal ViT class
mapped_state_dict = map_vit_state_dict(vit_ckpt)
tmm_albef_vision_encoder.load_state_dict(mapped_state_dict)

[XinhaoMei]

Thank you for your quick reply!
Understood, I will have a try.

Many thanks!