原帖在 Windows 启动 DataLoader worker 时失败,错误点名 init 内定义的 lambda;Linux 上的成功不能证明它可被 spawn 序列化:
AttributeError: Can't pickle local object 'NotMNISTDataLoader.__init__.<locals>.<lambda>'
函数可调用,不等于能被 spawn 子进程导入
原题在 Windows worker 启动阶段序列化局部 lambda 失败,尚未真正取图。Linux 默认启动方式可能掩盖这一点。应把变换函数写进可导入模块的顶层,而不是在 init 中另定义一个同名局部函数。
worker 启动需要序列化数据集与变换。把 Lambda 包装里的匿名函数换成模块顶层的具名函数,并在可导入的 .py 文件中运行。
import pickle
import torch
def repeat_to_rgb(tensor):
if tensor.ndim != 3 or tensor.shape[0] != 1:
raise ValueError("期望单张单通道 CHW 张量")
return tensor.repeat(3, 1, 1)
print("josh050608")
print("callable:", repeat_to_rgb.__module__, repeat_to_rgb.__qualname__)
print("pickle bytes:", len(pickle.dumps(repeat_to_rgb)))
在原变换序列中,用 transforms.Lambda(repeat_to_rgb) 代替局部 lambda,位置在 ToTensor 之后,因为该函数接收张量。函数不能再嵌套定义在 init 或另一个函数里。
如果源图片是 PIL,也可在 ToTensor 前使用 transforms.Grayscale(num_output_channels=3),避免这个自定义函数;是否需要保留颜色信息要按数据任务决定。
在 Windows 脚本中,把 DataLoader 创建和训练入口放到 if __name__ == '__main__': 内,Dataset 和变换函数定义留在模块顶层。Notebook 的临时函数不一定能被子进程按模块名导入。
先用 num_workers=0 验证变换类型与输出 [3,H,W],再使用 1 个 worker 测试,最后恢复需要的并行数。单进程成功只是绕过序列化,不能证明多进程问题已解决。
用一个独立脚本验证序列化与 batch
下面保存为 worker_probe.py,从终端执行。它使用合成灰度 Tensor 和一个 worker,不依赖 torchvision;类与函数都在模块顶层:
import torch
from torch.utils.data import Dataset, DataLoader
def to_three_channels(image):
if image.shape[0] != 1:
raise ValueError("此示例要求单通道输入")
return image.repeat(3, 1, 1)
class GrayDataset(Dataset):
def __len__(self):
return 2
def __getitem__(self, index):
return to_three_channels(torch.full((1, 4, 5), float(index))), index
if __name__ == "__main__":
loader = DataLoader(GrayDataset(), batch_size=2, num_workers=1,
multiprocessing_context="spawn")
images, labels = next(iter(loader))
assert images.shape == (2, 3, 4, 5)
assert labels.tolist() == [0, 1]
print("spawn batch:", images.shape)
原代码使用 transforms.Lambda 时,传入前文的顶层 repeat_to_rgb,放在 ToTensor 之后;若选择 PIL 阶段的现成灰度转三通道变换,则按该版本接口组合,不来回混用。
直接 pickle.dumps 成功只验证当前进程能生成引用,不完全证明新子进程能导入它。独立脚本中的 spawn 测试更接近故障条件;Windows、Notebook 与 Linux 默认启动方式仍应分别说明验证范围。
先 num_workers=0 验证样本内容,再显式 spawn 验证一个 worker,最后恢复所需并行数。单进程绕过序列化后能训练,不能宣称局部 lambda 的可导入性已经修复。
参考:原始求助帖;技术依据:Python multiprocessing、PyTorch 数据加载(2.5)、torchvision 0.8 变换文档。