How to use Avg 2d Pooling as a dataset transform?

I have a dataset of 2D images (or alternatively, 1 channel images) and I’d like to apply average 2D pooling as a transform. How can one do this? The following does not work:

    omniglot_dataset = torchvision.datasets.Omniglot(
        root=data_dir,
        download=True,
        transform=torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.CenterCrop((80, 80)),
            # torchvision.transforms.Resize((10, 10))
            torch.nn.functional.avg_pool2d(kernel_size=3, strides=1),
        ])
    )

You can use transforms.Lambda to call the functional API:

transform=torchvision.transforms.Compose([
    torchvision.transforms.CenterCrop((80, 80)),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Lambda(lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=3, stride=1)),
])

img = transforms.ToPILImage()(torch.randn(3, 224, 224))
out = transform(img)
1 Like