PyTorch加载自定义数据集训练GAN的技术求助

核心思路 #

PyTorch中所有内置数据集（包括MNIST）都继承自torch.utils.data.Dataset类，你只需要实现自己的数据集类继承该基类，就能无缝替换MNIST的代码逻辑。针对你的数据集结构，我们可以直接构建一个支持按训练/验证集和类别A/B加载的自定义数据集。

1. 实现自定义Dataset类 #

这个类会负责读取指定目录下的图像，并处理预处理逻辑：

from torch.utils.data import Dataset
from PIL import Image
import os

class CustomImageDataset(Dataset):
    def __init__(self, root_dir, split='Training', domain='A', transform=None):
        # 拼接目标文件夹路径：SET/Training/A 或 SET/Validation/B等
        self.image_dir = os.path.join(root_dir, split, domain)
        self.transform = transform
        # 收集所有图像文件路径（过滤非图像格式）
        self.image_paths = [
            os.path.join(self.image_dir, fname) 
            for fname in os.listdir(self.image_dir)
            if fname.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp'))
        ]

    def __len__(self):
        # 返回数据集样本总数
        return len(self.image_paths)

    def __getitem__(self, idx):
        # 加载单张图像并应用预处理
        img_path = self.image_paths[idx]
        # 统一转RGB避免灰度图和彩色图通道数不一致问题
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image

2. 构建数据加载器（DataLoader） #

和MNIST的用法完全一致，先定义图像预处理流程，再实例化数据集和加载器：

from torchvision import transforms
from torch.utils.data import DataLoader

# 定义预处理（根据你的GAN输入尺寸调整）
transform = transforms.Compose([
    transforms.Resize((64, 64)),  # 统一图像尺寸
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 归一化到[-1,1]，适配GAN常用的Tanh激活
])

# 实例化训练集和验证集
train_set_A = CustomImageDataset(root_dir='SET', split='Training', domain='A', transform=transform)
train_set_B = CustomImageDataset(root_dir='SET', split='Training', domain='B', transform=transform)
val_set_A = CustomImageDataset(root_dir='SET', split='Validation', domain='A', transform=transform)
val_set_B = CustomImageDataset(root_dir='SET', split='Validation', domain='B', transform=transform)

# 构建批量加载器
train_loader_A = DataLoader(train_set_A, batch_size=32, shuffle=True, num_workers=4)
train_loader_B = DataLoader(train_set_B, batch_size=32, shuffle=True, num_workers=4)

3. 适配GAN训练逻辑 #

原来用MNIST的训练循环可以直接复用，只需要把从MNIST加载的batch替换成自定义数据集的batch即可：

# 示例训练循环（核心逻辑和MNIST版本一致）
for epoch in range(100):
    # 遍历A和B的训练数据
    for imgs_A, imgs_B in zip(train_loader_A, train_loader_B):
        # 1. 用imgs_A作为输入，让生成器生成模仿B风格的图像
        # 2. 用真实的imgs_B和生成的图像训练判别器
        # 3. 更新生成器和判别器的参数
        # 所有逻辑和MNIST训练完全相同，只是输入张量是自定义图像的张量
        pass

注意事项 #

• 如果A和B的样本数量不一致，可以用itertools.cycle循环遍历数量较少的数据集，保证每个epoch都能配对：

  from itertools import cycle
  for imgs_A, imgs_B in zip(train_loader_A, cycle(train_loader_B)):
      # 训练逻辑
      pass
  

• 若你的图像是灰度图，可将convert('RGB')改为convert('L')，同时调整归一化的通道参数。

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