Torchvision.transforms implementation of Flatten()

Question

I have grayscale images, but I need transform it to a dataset of 1d vectors How can I do this? I could not find a suitable method in transforms:

train_dataset = torchvision.datasets.ImageFolder(root='./data',train=True, transform=transforms.ToTensor())
test_dataset = torchvision.datasets.ImageFolder(root='./data',train=False, transform=transforms.ToTensor())

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=4, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=4, shuffle=False)

Does this answer your question? [How do I flatten a tensor in pytorch?](https://stackoverflow.com/questions/55546873/how-do-i-flatten-a-tensor-in-pytorch) — Dishin H Goyani, Mar 28 '20 at 15:44
@Berriel Thank you, but not really. transforms.ToTensor returns Tensor, but I can't write in ImageFolder function 'transform = torch.flatten(transforms.ToTensor())' and it 'transform=transforms.LinearTransformation(transforms.ToTensor(),torch.zeros(1,784))' Maybe, it solved by transforms.Compose, but I don't know how — ComboSmash, Mar 28 '20 at 16:17
Check my answer. You used `torch.flatten` in the wrong way. You "have" to use Compose and I used it in my answer. — Berriel, Mar 28 '20 at 16:26

score 11 · Accepted Answer · answered Mar 28 '20 at 15:56

11

Here's how you can do it using Lambda

import torch
from torchvision.datasets import MNIST
import torchvision.transforms as T

# without flatten
dataset = MNIST(root='.', download=True, transform=T.ToTensor())
print(dataset[0][0].shape)
# >>> torch.Size([1, 28, 28])

# with flatten (using Lambda, but you can do it in many other ways)
dataset_flatten = MNIST(root='.', download=True, transform=T.Compose([T.ToTensor(), T.Lambda(lambda x: torch.flatten(x))]))
print(dataset_flatten[0][0].shape)
# >>> torch.Size([784])

answered Mar 28 '20 at 15:56

Berriel

12,659
4
43
67

Is it a good practice to do this operation in the transform section? – Opsse Jan 11 '21 at 17:33
1

@Opsse I wouldn't say good practice, but I'd rather consider this op part of my data processing pipeline than of my model; therefore I'd use it there. – Berriel Jan 11 '21 at 17:52

score 2 · Answer 2 · answered May 15 '22 at 08:48

The lambda seems to be unnecessary, and will raise a PyLint unnecessary-lambda / W0108 warning.

This version of @Berriel's solution is thus more precise:

import torch
from torchvision.datasets import MNIST
import torchvision.transforms as T

# without flatten
dataset = MNIST(root='.', download=True, transform=T.ToTensor())
print(dataset[0][0].shape)
# >>> torch.Size([1, 28, 28])

# with flatten (using Lambda, but you can do it in many other ways)
dataset_flatten = MNIST(root='.', download=True,
    transform=T.Compose([T.ToTensor(), T.Lambda(torch.flatten)]))
print(dataset_flatten[0][0].shape)
# >>> torch.Size([784])

I think even `T.Lambda(torch.flatten)` is unnecessary, and you can just replace it with `torch.flatten`, IE `dataset_flatten = torchvision.datasets.MNIST('./data', download=True, transform=torchvision.transforms.Compose( [torchvision.transforms.ToTensor(), torch.flatten]))` ? Works for me at least, Python 3.7.6 and PyTorch version 1.13.0+cu117 — Jake Levi, Apr 26 '23 at 09:57

Torchvision.transforms implementation of Flatten()

2 Answers2