手动加载ViT-H-14预训练权重遇state_dict尺寸不匹配问题求助

问题原因 #

你遇到的size mismatch for encoder.pos_embedding错误，核心原因是手动加载的权重文件与默认初始化的模型输入分辨率不匹配：

• 你下载的vit_h_14_swag-80465313.pth是针对518x518输入分辨率训练的SWAG权重，对应的位置嵌入包含1个CLS token + 37x37个patch嵌入（共1370个），形状为[1, 1370, 1280]。
• 直接调用torchvision.models.vit_h_14()初始化的模型，默认使用224x224输入分辨率，对应的位置嵌入是1个CLS token + 16x16个patch嵌入（共257个），形状为[1, 257, 1280]。

PyTorch官方的weights参数加载逻辑中，会自动对高分辨率的位置嵌入做插值适配，将其调整为模型默认输入尺寸对应的形状，因此不会触发错误。其他ViT模型用相同方法加载没问题，是因为你使用的权重文件刚好对应默认的224x224输入分辨率。

解决方案 #

方案1：初始化模型时匹配权重的输入分辨率 #

直接初始化适配518x518输入的模型，让位置嵌入形状与权重完全一致：

def vit_h_14(weight_path="/content/vit_h_14_swag-80465313.pth"):
    # 指定image_size为518，对应37x37个patch（518/14=37）
    pretrained_vit = torchvision.models.vit_h_14(image_size=518)
    pretrained_vit.load_state_dict(torch.load(weight_path))

    for parameter in pretrained_vit.parameters():
        parameter.requires_grad = False
    return pretrained_vit

同时需要修改预处理变换，将图像resize到518x518：

pretrained_transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((518, 518)),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

方案2：手动插值调整位置嵌入，适配默认224x224输入 #

如果想继续使用默认的224x224输入，可以手动对权重中的位置嵌入做插值调整，将高分辨率的patch嵌入缩放到16x16：

def vit_h_14(weight_path="/content/vit_h_14_swag-80465313.pth"):
    pretrained_vit = torchvision.models.vit_h_14()
    checkpoint = torch.load(weight_path)
    
    # 分离CLS token嵌入和patch位置嵌入
    pos_emb = checkpoint['encoder.pos_embedding']
    cls_token_emb = pos_emb[:, 0:1, :]  # 保留CLS token的嵌入
    patch_emb = pos_emb[:, 1:, :]       # 提取patch的位置嵌入
    
    # 将patch嵌入从37x37调整为16x16
    patch_emb = patch_emb.reshape(1, 37, 37, 1280).permute(0, 3, 1, 2)  # 转为[N, C, H, W]格式
    resized_patch_emb = torch.nn.functional.interpolate(
        patch_emb,
        size=(16, 16),
        mode='bicubic',
        align_corners=False
    )
    resized_patch_emb = resized_patch_emb.permute(0, 2, 3, 1).reshape(1, 256, 1280)  # 转回原格式
    
    # 合并CLS token和调整后的patch嵌入
    checkpoint['encoder.pos_embedding'] = torch.cat([cls_token_emb, resized_patch_emb], dim=1)
    
    # 加载处理后的权重
    pretrained_vit.load_state_dict(checkpoint)

    for parameter in pretrained_vit.parameters():
        parameter.requires_grad = False
    return pretrained_vit

这种方法可以继续使用官方默认的224x224预处理变换。

内容的提问来源于stack exchange，提问作者M. Al lail Qadrillah