PyTorch-Faster R-CNN模型训练好了后进行预测出现size mismatch for head.cls_loc.weight/cls_loc.bias/weight/bias

一、问题描述

在大牛的一个讲解训练Faster R-CNN的B站视频上,我依他的步骤训练完了模型。

然后进行预测的时候,出现了以下错误:

RuntimeError: Error(s) in loading state_dict for FasterRCNN:
	size mismatch for head.cls_loc.weight: copying a param with shape torch.Size([40, 2048]) from checkpoint, the shape in current model is torch.Size([36, 2048]).
	size mismatch for head.cls_loc.bias: copying a param with shape torch.Size([40]) from checkpoint, the shape in current model is torch.Size([36]).
	size mismatch for head.score.weight: copying a param with shape torch.Size([10, 2048]) from checkpoint, the shape in current model is torch.Size([9, 2048]).
	size mismatch for head.score.bias: copying a param with shape torch.Size([10]) from checkpoint, the shape in current model is torch.Size([9]).

二、解决思路

显然这是模型参数和输入参数之间不匹配的问题,但是我不知道问题出在哪个参数上呀?怎么办呢?

于是我回到那条视频中,发现有条评论和我的问题如出一辙。

PyTorch-Faster R-CNN模型训练好了后进行预测出现size mismatch for head.cls_loc.weight/cls_loc.bias/weight/bias_第1张图片

经由他们的发言,我怀疑我也可能是因为num_classes的数目存在问题。

于是我回到pycharm中,搜索num_classes。发现这个参数的得出是基于classes.txt文件的,这个文件里装了所有类别的名称,一行代表一名称,然后根据多少行得到多少个class。

读取classes.txt文件的函数如下:

PyTorch-Faster R-CNN模型训练好了后进行预测出现size mismatch for head.cls_loc.weight/cls_loc.bias/weight/bias_第2张图片

在我改成下面——num_classes的数目加1,之后问题得到解决。

PyTorch-Faster R-CNN模型训练好了后进行预测出现size mismatch for head.cls_loc.weight/cls_loc.bias/weight/bias_第3张图片

 

你可能感兴趣的:(PyTorch,pytorch,cnn,深度学习)