size mismatch for head.weight: copying a param with shape torch.Size([1000, 768]) from checkpoint

问题描述

---

我在github上的vit分类模型（https://github.com/YINYIPENG-EN/vit_classification_pytorch）使用了自己的数据集，训练的时候发生报错。

RuntimeError: Error(s) in loading state_dict for VisionTransformer:
        size mismatch for head.weight: copying a param with shape torch.Size([1000, 768]) from checkpoint, the shape in current model is torch.Size([3, 768]).
        size mismatch for head.bias: copying a param with shape torch.Size([1000]) from checkpoint, the shape in current model is torch.Size([3]).

报错原因：训练模型时出现FC层不匹配，我的数据集是 3个类的，下载的预权重是ImageNet 1000个类的。

解决方法

# 修改前的代码  train.py
def train(opt):
    classes_path = 'weights/cls_classes.txt'
    pretrained = False
    val_split = 0.1
    class_names, num_classes = get_classes(classes_path)
    if opt.model == 'vit':
        model = vit(num_classes=num_classes, pretrained=pretrained)
    if not pretrained:
        weights_init(model)
    if opt.weight != '':
        print('Loading {} into state dict...'.format(opt.weight))
        device = torch.device('cuda' if opt.cuda else 'cpu')
        model_dict = model.state_dict()
        pretrained_dict = torch.load(opt.weight, map_location=device)
        pretrained_dict = {k: v for k, v in pretrained_dict.items() if pretrained_dict.keys() == model_dict.keys()}
        model_dict.update(pretrained_dict)
        model.load_state_dict(model_dict)
    model_train = model.train()

方法一：将fc层忽略

找到pth文件导入的地方（即ckpt），将全连接层的参数直接pop掉

修改后的代码
def train(opt):
    classes_path = 'weights/cls_classes.txt'
    pretrained = False
    val_split = 0.1
    class_names, num_classes = get_classes(classes_path)
    if opt.model == 'vit':
        model = vit(num_classes=num_classes, pretrained=pretrained)
    if not pretrained:
        weights_init(model)
    if opt.weight != '':
        print('Loading {} into state dict...'.format(opt.weight))
        device = torch.device('cuda' if opt.cuda else 'cpu')
        # vit-patch_16.pth 是我的权重文件，你们要改成你自己的
        ckpt = torch.load('weights/vit-patch_16.pth', map_location='cpu') 
        ckpt.pop('head.weight')
        ckpt.pop('head.bias')
        model_dict = model.state_dict()
        pretrained_dict = torch.load(opt.weight, map_location=device)
        pretrained_dict = {k: v for k, v in pretrained_dict.items() if pretrained_dict.keys() == model_dict.keys()}
        model_dict.update(pretrained_dict)
        model.load_state_dict(ckpt, strict=False)
    model_train = model.train()

方法二：减去参数不匹配的层

def train(opt):
    classes_path = 'weights/cls_classes.txt'
    pretrained = False
    val_split = 0.1
    class_names, num_classes = get_classes(classes_path)
    if opt.model == 'vit':
        model = vit(num_classes=num_classes, pretrained=pretrained)
    if not pretrained:
        weights_init(model)
    if opt.weight != '':
        print('Loading {} into state dict...'.format(opt.weight))
        device = torch.device('cuda' if opt.cuda else 'cpu')
        model_dict = model.state_dict()
        pretrained_dict = torch.load(opt.weight, map_location=device)
       # 减去参数不匹配的层，我的层为“head”
        pretrained_dict = {k: v for k, v in pretrained_dict.items() if (k in model_dict and 'head' not in k)}
        model_dict.update(pretrained_dict)
        model.load_state_dict(model_dict)
    model_train = model.train()