strict=False后还是报错RuntimeError: Error(s) in loading state_dict for Siamese:

1. 问题

我们知道通过

model.load_state_dict(state_dict, strict=False)

可以暂且忽略掉模型和参数文件中不匹配的参数,先将正常匹配的参数从文件中载入模型。
但在使用时遇到了这样一个报错:

RuntimeError: Error(s) in loading state_dict for Siamese:
	size mismatch for fully_connect1.weight: copying a param with shape torch.Size([512, 2048]) from checkpoint, the shape in current model is torch.Size([512, 4608]).
	size mismatch for fc3.0.weight: copying a param with shape torch.Size([512, 2048]) from checkpoint, the shape in current model is torch.Size([512, 4608]).

一开始感觉很奇怪,我已经写明strict=False了,不匹配参数的不管就是了,为什么还要给我报错。

2. 原因及解决方案

经过仔细打印模型的键和文件中的键进行比对,发现是这样的:strict=False可以保证模型中的键与文件中的键不匹配时暂且跳过不管,但是一旦模型中的键和文件中的键匹配上了,PyTorch就会尝试帮我们加载参数,就必须要求参数的尺寸相同,所以会有上述报错。

比如在我们需要将某个预训练的模型的最后的全连接层的输出的类别数替换为我们自己的数据集的类别数,再进行微调,有时会遇到上述情况。这时,我们知道全连接层的参数形状会是不匹配,比如我们加载 ImageNet 1K 1000分类的预训练模型,它的最后一层全连接的输出维度是1000,但如果我们自己的数据集是10分类,我们需要将最后一层全链接的输出维度改为10。但是由于键名相同,所以PyTorch还是尝试给我们加载,这时1000和10维度不匹配,就会导致报错。

解决方案就是我们将 .pth 模型文件读入后,将其中我们不需要的层(通常是最后的全连接层)的参数pop掉即可。

以我训练模型为例:

model_pth = r"pth/epoch_064-val_loss_0.026-val_acc_1.00000.pth"
input_shape = [112, 112]
model = siamese(input_shape)
model.load_state_dict(torch.load(model_pth, map_location=device), False)

直接这样加载会出错,就是上面的错误:

RuntimeError: Error(s) in loading state_dict for Siamese:
	size mismatch for fully_connect1.weight: copying a param with shape torch.Size([512, 2048]) from checkpoint, the shape in current model is torch.Size([512, 4608]).
	size mismatch for fc3.0.weight: copying a param with shape torch.Size([512, 2048]) from checkpoint, the shape in current model is torch.Size([512, 4608]).

我们将最后 pth 文件加载进来之后(即 ckpt) 中全连接层的参数直接pop掉,至于需要pop掉哪些键名,就是上面报错信息中提到了的,在这里就是 fully_connect1.weight和 head.bias

ckpt.pop('fully_connect1.weight')
ckpt.pop('fc3.0.weight')

之后在运行,就正常了

即缺失了fully_connect1.weight和 head.bias这两个参数,也是正常的,因为在自己的数据集上微调时,我们本就不需要这两个参数,并且已经将它们从模型文件字典 ckpt 中pop掉了。现在,模型全连接之前的层(通常即所谓的特征提取层)的参数已经正常加载了,接下来可以在自己的数据集上进行微调。

因为反正我们也不用这些参数,就直接把这个键值对从字典中pop掉,以免 PyTorch 在帮我们加载时试图加载这些维度不匹配,我们也不需要的参数。

你可能感兴趣的:(目标检测,yolov5模型修改,模型部署,人工智能)