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.

Practice

1. What is the main purpose of the Flatten layer in PyTorch?

easy

A. To convert multi-dimensional input into a 1D vector per sample

B. To increase the number of channels in the input

C. To reduce the batch size during training

D. To apply activation functions element-wise

Solution

1. Step 1: Understand the role of Flatten layer

   The Flatten layer reshapes input data from multiple dimensions into a single long vector for each example, keeping batch size unchanged.

2. Step 2: Compare options with this role

   Only To convert multi-dimensional input into a 1D vector per sample describes this behavior correctly. Other options describe unrelated operations.

3. Final Answer:

   To convert multi-dimensional input into a 1D vector per sample -> Option A

4. Quick Check:

   Flatten layer = reshape to 1D vector [OK]

Hint: Flatten means reshape to 1D vector per example [OK]

Common Mistakes:

• Thinking Flatten changes batch size
• Confusing Flatten with convolution or activation
• Assuming Flatten adds or removes channels

2. Which of the following is the correct way to add a Flatten layer in a PyTorch nn.Sequential model?

easy

A. nn.Flatten(dim=0)

B. nn.Flatten(input_shape=(1, 28, 28))

C. nn.Flatten(start_dim=1)

D. nn.Flatten(start_dim=0)

Solution

1. Step 1: Recall PyTorch Flatten syntax

   PyTorch's nn.Flatten takes optional arguments start_dim and end_dim. By default, start_dim=1 flattens all dimensions except batch.

2. Step 2: Evaluate options

   nn.Flatten(input_shape=(1, 28, 28)) is invalid syntax. nn.Flatten(dim=0) uses unexpected keyword argument 'dim'. nn.Flatten(start_dim=0) flattens starting at batch dim (0), which is incorrect. nn.Flatten(start_dim=1) correctly specifies start_dim=1.

3. Final Answer:

   nn.Flatten(start_dim=1) -> Option C

4. Quick Check:

   Flatten start_dim=1 keeps batch dim [OK]

Hint: Use nn.Flatten(start_dim=1) to keep batch size [OK]

Common Mistakes:

• Using start_dim=0 which flattens batch dimension
• Passing input_shape argument (not supported)
• Using invalid keyword arguments like 'dim'

3. What is the output shape after applying nn.Flatten() to a tensor of shape (16, 3, 28, 28)?

medium

A. (16, 3, 28, 28)

B. (3, 28, 28)

C. (16, 28, 28)

D. (16, 2352)

Solution

1. Step 1: Understand input tensor shape

   The input tensor has shape (batch=16, channels=3, height=28, width=28).

2. Step 2: Calculate flattened size per example

   Flatten keeps batch size (16) and flattens remaining dims: 3*28*28 = 2352.

3. Final Answer:

   (16, 2352) -> Option D

4. Quick Check:

   Flatten output shape = (batch, product of other dims) [OK]

Hint: Multiply all dims except batch for flattened size [OK]

Common Mistakes:

• Forgetting to keep batch size dimension
• Using original shape without flattening
• Dropping batch dimension by mistake

4. Given the code below, what is the error and how to fix it?

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Conv2d(1, 10, kernel_size=3),
    nn.Flatten(start_dim=0),
    nn.Linear(10*26*26, 100)
)

medium

A. Conv2d output channels must match Linear input features

B. Flatten start_dim=0 flattens batch dimension; use start_dim=1 instead

C. Linear input size is incorrect; should be 10*28*28

D. Missing activation function after Conv2d

Solution

1. Step 1: Identify Flatten usage error

   Using start_dim=0 flattens batch dimension, which breaks batch processing.

2. Step 2: Correct Flatten start_dim

   Change start_dim=0 to start_dim=1 to keep batch size intact and flatten only feature dims.

3. Final Answer:

   Flatten start_dim=0 flattens batch dimension; use start_dim=1 instead -> Option B

4. Quick Check:

   Flatten start_dim=1 keeps batch size [OK]

Hint: Never flatten batch dimension; start_dim=1 keeps batch [OK]

Common Mistakes:

• Setting start_dim=0 flattens batch dimension
• Ignoring shape mismatch errors in Linear layer
• Assuming activation functions fix shape errors

5. You have a batch of images with shape (32, 3, 64, 64). You want to connect a convolutional network to a fully connected layer. Which PyTorch code correctly flattens the output before the dense layer?

hard

A. nn.Sequential(nn.Conv2d(3, 16, 3), nn.Flatten(start_dim=1), nn.Linear(16*62*62, 128))

B. nn.Sequential(nn.Conv2d(3, 16, 3), nn.Flatten(start_dim=0), nn.Linear(16*62*62, 128))

C. nn.Sequential(nn.Conv2d(3, 16, 3), nn.Flatten(), nn.Linear(3*64*64, 128))

D. nn.Sequential(nn.Conv2d(3, 16, 3), nn.Flatten(start_dim=1), nn.Linear(3*64*64, 128))

Solution

1. Step 1: Calculate output shape after Conv2d

   Conv2d with kernel_size=3 reduces each spatial dim by 2: 64 -> 62. Output shape: (32, 16, 62, 62).

2. Step 2: Flatten correctly and match Linear input

   Flatten with start_dim=1 keeps batch size 32 and flattens (16*62*62). Linear input features must match this product.

3. Final Answer:

   nn.Sequential(nn.Conv2d(3, 16, 3), nn.Flatten(start_dim=1), nn.Linear(16*62*62, 128)) -> Option A

4. Quick Check:

   Flatten start_dim=1 + correct Linear input size [OK]

Hint: Calculate Conv output size, flatten from dim=1, match Linear input [OK]

Common Mistakes:

• Flattening batch dimension (start_dim=0)
• Using wrong Linear input size
• Assuming default flatten matches input shape