PytorchのData Augmentationではまった話

PytorchのDataAugumentationで大いにハマりました。

torchvisionのtransforms.ImageFolderで以下のように変換処理を書いてトレーニングをしようとしていました。

    data_transform = {
        'train': transforms.Compose([
            transforms.RandomResizedCrop(
                image_size, scale=(0.5, 1.0)
            ),
            transforms.RandomHorizontalFlip(),
            transforms.RandomRotation(degrees=[-15, 15]),
            transforms.RandomErasing(0.5),
            transforms.ToTensor(),
            transforms.Normalize(mean, std),
        ]),
        'val': transforms.Compose([
            transforms.Resize(image_size),
            transforms.CenterCrop(image_size),
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ])
    }

    train_dataset = torchvision.datasets.ImageFolder(root=train_image_dir, transform=data_transform['train'], loader=myloader)
    val_dataset = torchvision.datasets.ImageFolder(root=val_image_dir, transform=data_transform['val'], loader=myloader)

ただ、トレーニングを始めようとすると以下のエラーを出します。

AttributeError: shape. Did you mean: 'save'?

原因は、transforms.RandomErasingの場所でした。
Pytorchの公式ドキュメントを読むと

Randomly selects a rectangle region in an torch Tensor image and erases its pixels. This transform does not support PIL Image.

とあり、RandomErasingはPILimageには適応できないと書いてありました。
なので、ToTensorでTensorに変換した後に記述しないとうまく読み取ってくれないようです。

以下のように直すとうまく動きました！

    data_transform = {
        'train': transforms.Compose([
            transforms.RandomResizedCrop(
                image_size, scale=(0.5, 1.0)
            ),
            transforms.RandomHorizontalFlip(),
            transforms.RandomRotation(degrees=[-15, 15]),
            transforms.ToTensor(),
            transforms.Normalize(mean, std),
            transforms.RandomErasing(0.5),
        ]),
        'val': transforms.Compose([
            transforms.Resize(image_size),
            transforms.CenterCrop(image_size),
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ])
    }

    train_dataset = torchvision.datasets.ImageFolder(root=train_image_dir, transform=data_transform['train'], loader=myloader)
    val_dataset = torchvision.datasets.ImageFolder(root=val_image_dir, transform=data_transform['val'], loader=myloader)

結局、データセットが悪いのか、モデルが悪いのか、、、学習はうまくいきませんでしたが、、、、