解决PyTorch Tensor尺寸不匹配运行时错误 | 深度学习实践

2025-01-27 02:29:37

解决Tensor尺寸不匹配的运行时错误

当在使用深度学习框架，特别是PyTorch时，我们有时会遇到类似于 RuntimeError: The size of tensor a (1024) must match the size of tensor b (640) at non-singleton dimension 1 这样的错误。这个错误表明，在某个操作中，两个张量的特定维度上的尺寸不匹配。在上面的例子中，这意味着 tensor a 的第 1 维的尺寸是 1024，而 tensor b 的第 1 维的尺寸是 640，但它们在计算中需要尺寸相同。这种错误通常发生在模型配置、权重加载或数据预处理出现问题时。下面将从原因分析、具体解决方法和安全建议三个方面进行阐述。

原因分析

根本原因在于模型初始化时，预训练模型（例如FastRCNNOutputLayers 中的分类层）的权重形状，和新数据集所需分类的数量不一致。具体来讲，上述报错信息中的 1024 指的是预训练模型中 cls_score 层权重的尺寸，例如 COCO 数据集中 80 个类别加一个背景类，再乘以 embedding 的维度，而 640 指的则是在当前修改的训练配置中，当前数据集中计算的维度。模型参数是基于预训练模型的，其分类器的输出维度已经预定义。当我们切换数据集并更改了类别数量(NUM_CLASSES)时，需要相应的调整模型中分类层（通常是全连接层）的维度，否则就产生了不兼容的权重。尤其是在使用预训练模型时，这点变得更加重要。

在给出的错误堆栈信息中可以看到，错误发生在 detectron2/modeling/roi_heads/fast_rcnn.py 中的 __init__ 方法的 self.cls_score.weight.copy_(pre_computed_w) 语句上。这表示尝试加载一个预先计算好的权重 pre_computed_w (来自于预训练模型)，但该权重的尺寸与当前模型的分类器权重 self.cls_score.weight 不兼容。

解决方案

1. 初始化分类层（cls_score）权重

当你的自定义数据集的类别数与预训练模型的不同时，应该放弃直接复制预训练权重，而是采用适当的初始化策略。一种方法是，对于分类层 cls_score 的权重，我们可以使用随机初始化。这样就避免了不匹配错误，同时也适配了新的分类类别数。修改 detectron2/modeling/roi_heads/fast_rcnn.py 文件中的相关代码:

    def __init__(self, cfg, input_shape):
        ...
        num_classes = cfg.MODEL.ROI_HEADS.NUM_CLASSES
        ...

        if self.use_sigmoid_ce:
            self.cls_score = nn.Linear(in_features, num_classes)
        else:
            self.cls_score = nn.Linear(in_features, num_classes + 1) # add one class for background

        # 不应该从预训练模型拷贝，应该采用随机初始化
        # if pre_computed_w is not None:
        #    self.cls_score.weight.copy_(pre_computed_w) # 不再直接拷贝预训练权重
        #    if pre_computed_b is not None:
        #       self.cls_score.bias.copy_(pre_computed_b)
        nn.init.normal_(self.cls_score.weight, std=0.01)
        nn.init.constant_(self.cls_score.bias, 0)
       
        self.box_pred = nn.Linear(in_features, num_classes * 4 if self.use_sigmoid_ce else (num_classes+1) * 4)
        nn.init.normal_(self.box_pred.weight, std=0.001)
        nn.init.constant_(self.box_pred.bias, 0)

操作步骤:

定位到 detectron2/modeling/roi_heads/fast_rcnn.py 文件。
修改 __init__ 方法中的 self.cls_score.weight.copy_(pre_computed_w) 部分。注释掉原始的拷贝代码，增加新的 nn.init.normal_ 和 nn.init.constant_初始化权重。
同样的需要初始化 self.box_pred.weight, 以避免相同的问题，加上nn.init.normal_(self.box_pred.weight, std=0.001)和nn.init.constant_(self.box_pred.bias, 0)
重新运行训练脚本。

此修改避免了直接拷贝权重导致的不匹配问题，为新的数据集初始化分类器参数，允许网络开始训练。
建议，在实践中, 对不同的参数应用不同的初始化方法可能会取得更好的效果。比如，可以将模型中的全连接层的bias 初始化为0。

2. 自定义权重加载逻辑

如果仍想复用部分预训练权重，可尝试根据类别数修改预训练模型的权重张量，使其与新数据集兼容，或者在加载预训练权重的时候有条件的修改其尺寸。在detectron2框架中，可以创建一个函数来处理预训练权重和当前模型参数的维度匹配，这个过程也可以是选择性的：对部分参数复用，对需要调整维度的参数则采用随机初始化。代码示例如下：

def load_pretrained_weights(model, pretrain_dict, num_classes):
    model_dict = model.state_dict()
    new_dict = {}
    for k,v in pretrain_dict.items():
      if k in model_dict:
        if "cls_score.weight" in k and v.shape[1]!= model_dict[k].shape[1]: # 这里以cls_score为例, 对于需要调整的参数,可以添加到if的条件里
          print("load custom cls weight!")
          new_v = torch.randn_like(model_dict[k]) * 0.01  # 使用随机初始化来调整
        else:
           new_v = v 
        new_dict[k] = new_v

    # 忽略不在模型结构中的key，和已经调整过维度的部分，并只加载部分权重
    update_dict = {k: v for k, v in new_dict.items() if k in model_dict and  model_dict[k].shape == new_dict[k].shape} 
    model_dict.update(update_dict) # 剩余没有修改的键值，以及形状完全匹配的，从预训练权重中拷贝过来
    model.load_state_dict(model_dict)

...
#在构建模型后调用该函数进行权重的加载。 
pre_model = torch.load(cfg.MODEL.WEIGHTS)["model"] # 加载权重
load_pretrained_weights(model, pre_model,cfg.MODEL.ROI_HEADS.NUM_CLASSES) 
...

操作步骤:

定义load_pretrained_weights函数，并指定你需要特殊处理的参数
在 detectron2/engine/defaults.py 或者 tools/train_net.py 中找到权重加载的部分。
将预训练模型读取进来, 修改默认的模型权重读取方法为使用你自定义的load_pretrained_weights方法，该函数在 detectron2 模型建立后调用。
重新运行训练脚本。

这种方法灵活性更高，可以让你控制如何以及在何种程度上利用预训练模型的参数。
建议权重调整的维度逻辑根据你的具体模型进行调整，注意要保证所有参数初始化到合适的大小。

其他注意事项

数据集类别数配置 确认配置文件（通常是yaml文件）中MODEL.ROI_HEADS.NUM_CLASSES 与自定义数据集类别数目一致。
模型结构的兼容性 确保模型结构中关键层（例如 cls_score层，也就是全连接层）的输入维度参数符合逻辑。特别是在更改了backbone的时候需要特别注意。
错误追踪 阅读完整的报错信息可以帮助确定问题代码位置，缩小排错范围。

通过理解错误的根本原因，并采取上述的调整或权重加载策略，能够有效解决“Tensor尺寸不匹配” 的运行时错误。合理的初始化和参数处理对于成功训练深度学习模型是不可或缺的。