PyTorch自定义CNN数据集报错：TypeError及标签异常求助

错误原因拆解 #

1. TypeError：图像张量用作列表索引
   show_data函数中，labels[data_sample[0]]试图用形状为(1, 16, 16)的图像张量作为列表索引，而列表索引要求单个整数张量或整数，因此触发错误。
2. 标签显示错误
   将标题中的data_sample[0]改为data_sample[1]后，若labels列表顺序与ImageFolder自动生成的类别索引不匹配，就会出现标签显示错误。ImageFolder会按文件夹名称的字典序分配类别索引，需确保自定义labels列表与dataset.class_to_idx的顺序一致。
3. 'int' object has no attribute 'numpy'
   若同时修改plt.imshow中的data_sample[0]为data_sample[1]，则会尝试对整数标签调用numpy()方法，整数类型无此属性，导致报错。

修复后的完整代码 #

数据集构建代码（修正后） #

import torch
from torchvision import datasets
from torchvision import transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

IMAGE_SIZE = 16
# 加载数据集，调整变换顺序（先转灰度再转张量）
dataset = datasets.ImageFolder(
    '/content/gdrive/MyDrive/Dataset',
    transforms.Compose([
        transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
        transforms.Grayscale(),
        transforms.ToTensor()
    ])
)

# 拆分训练集和验证集
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
dataset_train, dataset_val = torch.utils.data.random_split(dataset, [train_size, test_size])

# 自定义标签，需与dataset.class_to_idx的顺序一致
labels = ["5","10","20","30","40","50","75","100","150","200"]
# 验证类别索引映射是否匹配
print("类别索引映射：", dataset.class_to_idx)
print("自定义标签顺序：", labels)

# 测试DataLoader（可选）
train_loader = DataLoader(dataset_train, batch_size=16, shuffle=True)
for batch in train_loader:
    print("图像张量形状：", batch[0].shape)
    print("标签张量：", batch[1])
    break

测试与可视化代码（修正后） #

import random

# 假设model已定义并训练完成
model.eval()

# 随机选取验证集样本
random_idx = random.randint(0, len(dataset_val) - 1)
test_sample = dataset_val[random_idx]

def show_data(data_sample, size):
    # data_sample[0]是图像张量，转换为numpy数组后显示
    img = data_sample[0].numpy().reshape(size, size)
    plt.imshow(img, cmap='gray')
    # data_sample[1]是整数标签，用于索引自定义labels列表
    true_label = labels[data_sample[1]]
    plt.title(f"真实标签 = {true_label}")

# 可视化第一个数据集样本
for n, sample in enumerate(dataset):
    show_data(sample, IMAGE_SIZE)
    plt.show()
    if n == 0:
        break

# 可视化验证集样本
show_data(test_sample, IMAGE_SIZE)
plt.show()

# 模型预测：调整张量形状并转换结果为整数
input_tensor = test_sample[0].view(-1, 1, IMAGE_SIZE, IMAGE_SIZE)
with torch.no_grad():
    output = model(input_tensor)
# 将张量类型的预测索引转为整数
pred_idx = torch.argmax(output, dim=1).item()
pred_label = labels[pred_idx]
print(f"AI预测结果：{pred_label}")

关键修复点说明 #

• 变换顺序调整：将transforms.Grayscale()移至transforms.ToTensor()之前，确保生成单通道的灰度图像张量。
• 索引明确区分：data_sample[0]对应图像张量，用于显示；data_sample[1]对应整数标签，用于索引labels列表。
• 类别映射验证：打印dataset.class_to_idx确认自动分配的索引与自定义labels顺序一致，避免标签错位。
• 预测结果转换：用item()方法将torch.argmax返回的张量转为整数，才能正确索引labels列表。

内容的提问来源于stack exchange，提问作者A.Serin