PyTorch Hub加载的I3D模型修改输出层报错，如何适配自定义数据集？

问题场景 #

想要将Torch Hub中预训练在Kinetics 400数据集上的I3D模型（i3d_r50）微调至包含4个类别的自定义数据集，加载模型后尝试直接通过类名修改输出层时，出现AttributeError: 'Net' object has no attribute 'ResNetBasicHead'错误。

错误原因 #

PyTorchVideo中的i3d_r50模型并未将ResNetBasicHead设为模型的直接属性，该分类头是模型blocks模块列表的最后一个元素，无法通过类名直接访问。

正确修改代码 #

import torch
# 加载预训练模型
model = torch.hub.load("facebookresearch/pytorchvideo", 'i3d_r50', pretrained=True)
num_classes = 4

# 获取模型的分类头模块（ResNetBasicHead）
classification_head = model.blocks[-1]
# 替换输出层为4分类
classification_head.proj = torch.nn.Linear(classification_head.proj.in_features, num_classes)

验证方式 #

打印修改后的输出层结构，确认out_features已更新为4：

print(model.blocks[-1].proj)

内容的提问来源于stack exchange，提问作者Carlo