记录pytorch中DataLoader使用torchvision.transforms.Lambda进行图片变换时的问题

1.错误代码

import sys
sys.path.append('..')
import torch
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision import datasets,transforms
if __name__ == '__main__':

    transform = transforms.Compose([
        transforms.Scale(256),
        transforms.FiveCrop([224,224]),
        transforms.Lambda(lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])),
    ])

    dataset = ImageFolder(root='../data/UCF101/test/image',
                              transform=transform,
                              )

    dataloader = DataLoader(dataset, batch_size=4,
                               shuffle=True, num_workers=1, drop_last=True)

    for i_batch, img in enumerate(dataloader):
        print(type(img), i_batch)

2.产生的错误

ForkingPickler(file, protocol).dump(obj)
_pickle.PicklingError: Can't pickle  at 0x0000021DDB86AC80>: attribute lookup  on __main__ failed

3.问题的解决：

import sys
sys.path.append('..')
import torch
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision import datasets,transforms

def totensor(crops):
    return torch.stack([transforms.ToTensor()(crop) for crop in crops])

if __name__ == '__main__':

    transform = transforms.Compose([
        transforms.Scale(256),
        transforms.FiveCrop([224,224]),
        transforms.Lambda(totensor),
    ])

    dataset = ImageFolder(root='../data/UCF101/test/image',
                              transform=transform,
                              )

    dataloader = DataLoader(dataset, batch_size=4,
                               shuffle=True, num_workers=1, drop_last=True)

    for i_batch, img in enumerate(dataloader):
        print(type(img), i_batch)

4.总结

根据这个错误，在网上搜索，大多的解决方案，就是将lambda方法放在外面实现，事实证明这确实可行。希望对其他人有帮助吧~