PyTorch CNN单张图像预测出现索引越界错误如何解决

你遇到的索引越界问题核心是模型输出维度和实际分类数量不匹配：

• 你搭建的BrainTumorClassification模型最后一层全连接层设置为nn.Linear(512,6)，意味着模型会输出6个类别的置信度，经过argmax计算得到的索引取值范围是0~5
• 你的分类标签列表classes只有4个元素，合法索引范围是0~3，当模型预测输出的索引是4或5时，访问classes[index]就会触发越界错误

按以下步骤操作即可解决：

1. 修正模型结构，将最后一层全连接层的输出维度调整为和分类数一致的4：

# 修改前代码
nn.Linear(512,6)
# 修改后代码
nn.Linear(512,4)

2. 重新训练模型，修改模型结构后原有训练的权重维度和新结构不匹配，必须重新训练得到适配4分类的权重后再进行预测
3. 可选补充：可以在预测函数中增加索引合法性校验，避免极端情况报错：

def classify(image_path,image_transforms, classes):
    image = Image.open(image_path)
    image_tensor = image_transforms(image).float()
    image_tensor = image_tensor.unsqueeze_(0)
    input = Variable(image_tensor)
    output = model(input)
    index = output.data.numpy().argmax()
    # 新增合法性校验
    if 0 <= index < len(classes):
        pred = classes[index]
    else:
        pred = "未知类别"
    return pred

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