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

一、问题描述

在大牛的一个讲解训练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]).

二、解决思路

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

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

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

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

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

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