使用PyTorch DataLoader加载数据时如何获取文件名与文件路径

你当前使用PyTorch 1.8 + Python 3.8通过torchvision.datasets.ImageFolder读取本地目录下的图像数据，需要在遍历DataLoader时同时获取对应图像的文件路径、文件名，可通过以下两种方式实现：

方法一：自定义继承ImageFolder的数据集类（推荐） #

PyTorch原生的ImageFolder会自动把所有样本的路径、对应类别ID存在samples属性中，仅需重写__getitem__方法将路径加入返回值即可，该方法兼容shuffle、多进程加载、drop_last等所有DataLoader参数设置，不会出现路径和图像匹配错误的问题。

1. 首先定义自定义数据集类：

from torchvision.datasets import ImageFolder

class ImageFolderWithPaths(ImageFolder):
    # 重写方法额外返回图像文件路径
    def __getitem__(self, index):
        img, label = super().__getitem__(index)
        # 取当前索引对应样本的路径
        path = self.samples[index][0]
        return img, label, path

2. 替换原有数据集定义代码：

注意修正你原有代码中变量名不匹配的问题：之前定义的训练变换变量为transform_train，你写的train_transforms属于笔误，测试集同理

# 替换原来的datasets.ImageFolder为自定义类即可
train_data = ImageFolderWithPaths(data_dir + '/train', transform = transform_train)
test_data = ImageFolderWithPaths(data_dir + '/test', transform = transform_test)

# DataLoader定义不需要修改
trainloader = torch.utils.data.DataLoader(train_data, batch_size = 32, shuffle=True)
testloader = torch.utils.data.DataLoader(test_data, batch_size = 32)

3. 遍历直接获取路径：

import os

for imgs, labs, paths in trainloader:
    print(f"图像张量形状：{imgs.shape}，标签形状：{labs.shape}")
    # 取批次内第一个样本的路径和文件名
    first_sample_path = paths[0]
    first_sample_name = os.path.basename(first_sample_path)
    print(f"第一个样本路径：{first_sample_path}，文件名：{first_sample_name}")
    break

方法二：通过样本索引匹配路径（临时使用） #

如果不想修改原有数据集定义，可以直接调用ImageFolder的samples属性拿到全量路径列表，通过批次索引匹配对应路径，仅需注意如果DataLoader开启了shuffle，要固定随机种子保证索引对应正确。

# 提前获取全量训练样本路径列表
train_all_paths = [sample[0] for sample in train_data.samples]
batch_size = trainloader.batch_size

for batch_idx, (imgs, labs) in enumerate(trainloader):
    # 计算当前批次对应的样本索引范围
    start_idx = batch_idx * batch_size
    end_idx = start_idx + imgs.shape[0]
    current_batch_paths = train_all_paths[start_idx: end_idx]
    print(f"当前批次共{len(current_batch_paths)}个样本的路径已获取")
    break

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