num_workers>0时Image Tensors值全为0的技术求助

核心问题分析 #

你的问题本质是多进程环境下预加载的张量数据没有正确传递到子进程，导致子进程读取到的张量数值全为0。单进程模式下数据在主进程直接读取，没有序列化/跨进程传递的问题；而num_workers>0时，PyTorch会启动子进程加载数据，当你把预加载的大张量直接作为Dataset的属性时，spawn启动方式会对整个Dataset进行序列化复制，在WSL+Jupyter的交互式环境下，容易出现张量数据丢失或未正确反序列化的情况。

之前添加的pickle/dill相关代码解决了序列化报错，但并没有解决张量跨进程传递的核心问题。

具体解决办法 #

办法1：修改Dataset，从文件按需加载（最推荐） #

放弃预加载所有张量到内存，改为在__getitem__中从.pt文件加载单个样本，彻底避免跨进程传递大张量的问题。

修改后的Dataset代码：

class SRDataset(Dataset):
    def __init__(self, hr_paths, lr_paths, transform=None):
        self.hr_paths = hr_paths  # 存储.pt文件路径列表，而非预加载的张量
        self.lr_paths = lr_paths
        self.transform = transform

        assert len(self.hr_paths) == len(self.lr_paths), \
            "Number of HR and LR images must be equal"

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

    def __getitem__(self, idx):
        # 按需加载单个张量
        hr_img = torch.load(self.hr_paths[idx])
        lr_img = torch.load(self.lr_paths[idx])

        if self.transform:
            hr_img = self.transform(hr_img)
            lr_img = self.transform(lr_img)

        return {'lr': lr_img, 'hr': hr_img}

对应的数据集加载函数也要调整为传入文件路径列表：

def create_sr_datasets_and_loaders(hr_paths, lr_paths, batch_size=32, transform=None):
    full_dataset = SRDataset(hr_paths, lr_paths, transform)
    total_size = len(full_dataset)
    train_size = int(0.8 * total_size)
    val_size = int(0.1 * total_size)
    test_size = total_size - train_size - val_size

    train_dataset, val_dataset, test_dataset = random_split(
        full_dataset,
        [train_size, val_size, test_size],
        generator=torch.Generator().manual_seed(seed)
    )

    # 保持DataLoader设置不变
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=1,
        drop_last=True,
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=1,
        drop_last=True,
    )
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=1,
        drop_last=True,
    )
    return {
        'datasets': {
            'train': train_dataset,
            'val': val_dataset,
            'test': test_dataset
        },
        'loaders': {
            'train': train_loader,
            'val': val_loader,
            'test': test_loader
        }
    }

办法2：使用共享内存传递预加载张量（适合必须预加载的场景） #

如果必须提前把所有张量加载到内存，可以用PyTorch的共享内存机制，让子进程直接访问主进程的张量内存，避免复制：

在创建Dataset前，把张量移到共享内存：

# 预加载张量后，转为共享内存张量
hr_tensors = [t.share_memory_() for t in hr_tensors]
lr_tensors = [t.share_memory_() for t in lr_tensors]

# 然后再传入Dataset
data = create_sr_datasets_and_loaders(hr_tensors, lr_tensors, batch_size=64, transform=None)

同时，调整多进程启动策略的代码，确保在Jupyter中正确初始化（放在第一个单元格最顶部）：

import torch.multiprocessing as mp

try:
    mp.set_start_method('spawn', force=True)
except RuntimeError:
    pass
torch.multiprocessing.set_sharing_strategy("file_system")

办法3：禁用Jupyter的交互式输出干扰（辅助调试） #

Jupyter的交互式环境可能会干扰多进程的输出和内存管理，你可以尝试：

• 把数据加载和验证的代码放到一个单独的函数中，用%run魔法命令执行，而不是分步运行
• 确保所有多进程相关的代码在同一个单元格中执行，避免跨单元格的状态污染

环境适配注意事项 #

在WSL2+Jupyter环境下：

• spawn是唯一稳定的多进程启动方式，fork在WSL中容易出现各种内存问题
• 避免在Dataset中存储大体积的可序列化对象，尽量用按需加载的方式
• 如果使用共享内存，确保张量是CPU张量（GPU张量的共享内存机制在WSL中可能有兼容性问题）

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