Flatten layer in PyTorch - Practice Problems & Coding Challenges

Challenge - 5 Problems

❓ Predict Output

intermediate

Output shape after Flatten layer in PyTorch

Given the following PyTorch code, what is the shape of the tensor after applying the Flatten layer?

PyTorch

import torch
import torch.nn as nn

x = torch.randn(10, 3, 28, 28)  # batch of 10 images, 3 channels, 28x28 pixels
flatten = nn.Flatten()
out = flatten(x)
print(out.shape)

A torch.Size([10, 2352])

B torch.Size([10, 3, 28, 28])

C torch.Size([10, 28, 28, 3])

D torch.Size([10, 3, 784])

💡 Hint

Flatten combines all dimensions except the batch dimension into one.

The Flatten layer keeps the batch size (10) and flattens the remaining dimensions (3*28*28 = 2352) into one dimension.

❓ Model Choice

intermediate

Choosing Flatten layer position in a CNN

In a convolutional neural network (CNN), where should you place the Flatten layer?

A After the last convolutional or pooling layer and before the fully connected layers

B Before the first convolutional layer

C After the output layer

D Between two convolutional layers

💡 Hint

Flatten prepares data for fully connected layers by converting multi-dimensional data to 2D.

Flatten is used after convolutional and pooling layers to convert the multi-dimensional feature maps into a 2D tensor for fully connected layers.

❓ Hyperparameter

advanced

Effect of start_dim parameter in PyTorch Flatten

What is the effect of setting start_dim=2 in nn.Flatten(start_dim=2) when applied to a tensor of shape (5, 4, 3, 2)?

A The tensor shape remains unchanged

B The tensor is flattened starting from dimension 1, resulting in shape (5, 24)

C The tensor is flattened starting from dimension 0, resulting in shape (120,)

D The tensor is flattened starting from dimension 2, resulting in shape (5, 4, 6)

💡 Hint

start_dim defines the first dimension to flatten from.

With start_dim=2, dimensions 2 and 3 (3 and 2) are flattened into one dimension (3*2=6), so shape becomes (5, 4, 6).

🔧 Debug

advanced

Debugging Flatten layer usage error

What error will this PyTorch code raise and why? import torch import torch.nn as nn x = torch.randn(8, 3, 32, 32) flatten = nn.Flatten(start_dim=4) out = flatten(x)

A IndexError: start_dim must be less than or equal to the number of dimensions minus 1

B RuntimeError: input tensor must be 2D or more

C No error, output shape is (8, 3, 32, 32)

D TypeError: Flatten() got an unexpected keyword argument 'start_dim'

💡 Hint

start_dim can equal input.dim(); it flattens nothing.

The tensor has 4 dimensions. start_dim=4 equals input.dim(), so it flattens an empty range (dim 4 to 4), causing no change and no error.

🧠 Conceptual

expert

Why Flatten layer is crucial before fully connected layers

Why is the Flatten layer necessary before feeding data into fully connected (linear) layers in neural networks?

A Because Flatten normalizes the data to have zero mean and unit variance

B Because fully connected layers require 2D input: (batch_size, features), and Flatten converts multi-dimensional tensors to this shape

C Because Flatten reduces the number of parameters in the model

D Because Flatten applies non-linear activation to the input

💡 Hint

Think about the input shape requirements of linear layers.

Fully connected layers expect 2D input where each row is a sample and columns are features. Flatten reshapes multi-dimensional tensors to this format.