PyTorch Lambda类与普通lambda的差异及存在意义探讨

确实，普通lambda作为可调用对象能满足不少transform组合场景，但PyTorch提供torchvision.transforms.Lambda类主要是解决几个lambda搞不定的核心痛点：

• 序列化支持：普通lambda函数无法被pickle序列化，如果你需要保存包含transform链的数据集、模型或者训练状态（比如用torch.save），直接用lambda会抛出序列化错误。而Lambda类是可序列化的，能完美适配这种场景。
• 统一接口规范：和torchvision里其他transform组件（比如ToTensor、Resize）保持一致的类实例风格，让整个transform链的代码更统一、更易维护，尤其是在多人协作的项目里，大家能一眼识别这是一个标准的transform操作。
• 可读性与复用性：虽然lambda写起来简洁，但Lambda类可以让你更清晰地封装预处理逻辑，比如把复杂的操作包装成Lambda实例后，能方便地在多个transform链中复用，比到处复制lambda代码更整洁。
• 兼容与扩展便利：在一些早期PyTorch版本中，部分transform组合工具对纯lambda的支持有限，Lambda类的存在确保了兼容性；另外，如果你需要给transform添加额外的日志、调试逻辑，基于Lambda类扩展会比直接修改lambda更灵活。

举个序列化的实际例子对比：

# 使用普通lambda会触发序列化错误
from torchvision import transforms
from torch.utils.data import Dataset
import torch

class MyDataset(Dataset):
    def __init__(self):
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            lambda x: x * 2  # 该lambda无法被pickle序列化
        ])
    
    def __getitem__(self, idx):
        return self.transform(torch.randn(3, 224, 224))
    
    def __len__(self):
        return 10

# 尝试保存数据集会失败
dataset = MyDataset()
torch.save(dataset, "dataset.pth")  # 抛出 PicklingError

# 替换为Lambda类即可正常工作
class MyDataset(Dataset):
    def __init__(self):
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Lambda(lambda x: x * 2)
        ])
    
    def __getitem__(self, idx):
        return self.transform(torch.randn(3, 224, 224))
    
    def __len__(self):
        return 10

dataset = MyDataset()
torch.save(dataset, "dataset.pth")  # 保存成功

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