修改Resnet34的全连接层,并加入了Involution模块,为了做消融实验,就需要加载预训练权重,此时,原先训练好的权重文件的全连接层会与我模型的全连接层冲突,因此就不必加载训练好的全连接层权重。
model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) if pretrained: pretrained_dict = model_zoo.load_url(model_urls['resnet34']) model_dict = model.state_dict() pretrained_dict.pop('fc.weight') pretrained_dict.pop('fc.bias') # 筛除不加载的层结构 pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} # 更新当前网络的结构字典 model_dict.update(pretrained_dict) model.load_state_dict(model_dict)
修改Resnet的全连接层后加载预训练模型
最新推荐文章于 2024-08-30 08:59:25 发布