遍历PyTorch DataLoader时陷入无限循环问题求助

我正在创建一个包含10000个样本（8000个训练样本、2000个验证样本）的3D测试数据集，用于测试我的3D CNN模型。在使用next(iter(Train_dataloader))查看Train_dataloader的第一批数据时，程序陷入无限循环，内核始终无法停止。

自定义数据集及DataLoader创建代码如下：

class Binary3DDataset(Dataset):
    def __init__(self, data,transform=None):
        self.data = data
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data[idx]
        sample =self.transform(sample)
        return sample
 

# Define data augmentation
data_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomAffine(degrees = 15, translate = (0.1,0.1)),
    transforms.ToTensor(),
])

# create 10000 random 3D sample with dimension [1, 32,32,32] "[#of channel, depth, height,width]"
num_samples = 10000  
voxel_sample = np.random.choice([0, 1], size=(10000,1,32,32,32), p=[0.7, 0.3])

# Create an instance of custom dataset with data augmentation
augumented_custom_dataset = Binary3DDataset(voxel_sample,transform = data_transform)

# Create train and test dataloader to iterate over the augmented data
batch_size = 32
train_dataset, test_dataset = random_split(augumented_custom_dataset, [8000,2000])

Train_dataloader = DataLoader(train_dataset, batch_size = batch_size, shuffle = True,num_workers=1)
Test_dataloader = DataLoader(test_dataset,batch_size = batch_size, shuffle = False,num_workers=1)

调用以下代码时陷入无限循环：

# Get one batch from the train_dataloader
data_iter = iter(Train_dataloader)
inputs = next(data_iter)

已确认数据集和DataLoader长度正常：

len(train_dataset),len(test_dataset),len(Train_dataloader),len(Test_dataloader)

输出：(8000, 2000, 250, 63)

问题原因 #

你使用的torchvision.transforms中的大部分操作（比如RandomHorizontalFlip、RandomAffine）是针对2D图像设计的，无法正确处理3D张量（维度为[C, D, H, W]）：

• RandomAffine仅支持2D空间变换，传入3D张量时会导致内部逻辑异常，进而引发无限循环；
• 即使是RandomHorizontalFlip，默认也是对最后一个维度（宽度）翻转，但没有明确支持3D场景的深度维度处理，可能引发未定义行为；
• ToTensor的设计是将(H, W, C)格式的numpy数组转为(C, H, W)的张量，但你的输入已经是(C, D, H, W)的3D格式，会导致维度处理混乱。

解决方案 #

- 使用支持3D数据的增强库 #

推荐使用MONAI（专为医学3D图像设计的PyTorch工具库）的transforms，它原生支持3D张量操作：

import monai.transforms as transforms

data_transform = transforms.Compose([
    transforms.RandFlip(spatial_axis=2),  # 3D水平翻转（针对width维度）
    transforms.RandAffine(
        degrees=15, 
        translate=(0.1, 0.1, 0.1),  # 3D平移（D/H/W三个维度）
        spatial_size=(32,32,32)
    ),
    transforms.ToTensor(),  # MONAI的ToTensor支持3D格式
])

- 自定义3D版本的transforms #

如果不想引入额外库，可以自己实现3D兼容的变换：

import torch
import numpy as np

class RandomHorizontalFlip3D:
    def __init__(self, p=0.5):
        self.p = p
    
    def __call__(self, sample):
        # sample shape: [C, D, H, W]
        if np.random.rand() < self.p:
            sample = sample[:, :, :, ::-1]
        return sample

class RandomAffine3D:
    def __init__(self, degrees=15, translate=(0.1,0.1)):
        self.degrees = degrees
        self.translate = translate
    
    def __call__(self, sample):
        # 简化实现：仅对H-W平面做仿射变换，保持D维度不变
        from torchvision.transforms.functional import affine
        sample = torch.tensor(sample)
        angle = np.random.uniform(-self.degrees, self.degrees)
        tx = np.random.uniform(-self.translate[0], self.translate[0]) * sample.shape[2]
        ty = np.random.uniform(-self.translate[1], self.translate[1]) * sample.shape[3]
        # 对每个深度切片应用2D仿射
        augmented = torch.stack([affine(slice_, angle, (tx, ty), 1.0, 0.0) for slice_ in sample[0]])
        return augmented.unsqueeze(0)

data_transform = transforms.Compose([
    RandomHorizontalFlip3D(p=0.5),
    RandomAffine3D(degrees=15, translate=(0.1,0.1)),
    lambda x: torch.tensor(x, dtype=torch.float32),  # 替代ToTensor，直接转张量
])

- 调整数据格式与transform顺序（不推荐） #

如果坚持使用torchvision的transforms，需要先将3D张量转为2D切片处理，再重新堆叠，但这种方式效率较低：

class Convert3DTo2D:
    def __call__(self, sample):
        # [C, D, H, W] -> [D*C, H, W]
        return sample.reshape(-1, sample.shape[2], sample.shape[3])

class Convert2DTo3D:
    def __init__(self, original_shape):
        self.original_shape = original_shape
    
    def __call__(self, sample):
        # [D*C, H, W] -> [C, D, H, W]
        return sample.reshape(self.original_shape)

data_transform = transforms.Compose([
    Convert3DTo2D(),
    transforms.RandomHorizontalFlip(),
    transforms.RandomAffine(degrees=15, translate=(0.1,0.1)),
    transforms.ToTensor(),
    Convert2DTo3D((1,32,32,32)),
])

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