ConvNextv2预训练权重加载到微调模型

ICCV 2023 《ConvNeXt V2: Co-designing and Scaling ConvNets with Masked Autoencoders》

Github:https://github.com/facebookresearch/ConvNeXt-V2

针对ImageNet-1K fine-tuned models

    model = convnextv2_base(pretrained=False).cuda()
    model_dict = model.state_dict()
    
    pretrained_path = "预训练模型路径"
    pretrained_weight = torch.load(pretrained_path)['model']
     #fine-tuned模型的键是'model'

    pretrained_weight = {k: v for k, v in pretrained_weight.items() if k in model_dict}
    # 只保留预训练模型中，自己建的model有的参数
    model_dict.update(pretrained_weight)  # 将预训练的值，更新到自己模型的dict中
    model.load_state_dict(model_dict)  # model加载dict中的数据，更新网络的初始值
    
    num_fr = model.head.in_features
    model.head = nn.Linear(num_fr, classes)