Use nn.Flatten Layer

easyFREEModelingPyTorch ModelingPytorch

Problem Statement

Use nn.Flatten to convert multi-dimensional tensors to 1D (per batch).

Background

When connecting convolutional layers to fully connected layers, you need to flatten the spatial dimensions while keeping the batch dimension intact. nn.Flatten does this automatically.

Your Task

Write a function apply_flatten(batch_size, channels, height, width) that creates a random tensor with the given shape, flattens it, and returns information about the shapes before and after.

Output Format

Return a dictionary with keys: "input_shape" (list), "output_shape" (list), "flattened_features" (int).

Example

Input

batch_size=2, channels=3, height=4, width=4

Output

{'input_shape': [2, 3, 4, 4], 'output_shape': [2, 48], 'flattened_features': 48}

Reasoning

Inputs: batch_size=2, channels=3, height=4, width=4 describe a batch of 2 images, each with 3 channels on a $4 \times 4$ grid. The random tensor therefore has shape [2, 3, 4, 4], which is input_shape.
The math: flattening unrolls everything except the batch axis into one long vector. Each image holds $C \times H \times W = 3 \times 4 \times 4 = 48$ numbers, so each one becomes a vector of length 48.
Output shape: the batch axis is untouched, giving $2 \times 48$ , i.e. output_shape is [2, 48]. The total element count is unchanged: $2 \cdot 3 \cdot 4 \cdot 4 = 96 = 2 \cdot 48$ .
flattened_features is the per-sample length, 48. This is exactly the input size a fully connected layer placed after this point would need.
Notice: the random values never matter here, only the shapes. A $1 \times 28 \times 28$ MNIST image would flatten to $784$ features per sample, and a batch of 4 would stay 4 rows.

Constraints

Use nn.Flatten()
Batch dimension should be preserved
Use torch.randn for random tensor

Learn this in

PyTorch FundamentalsFoundations

More from Pytorch

Apply Softmax for Predictionseasy Count Model Parametersmedium Build with nn.Sequentialeasy Create a Simple Neural Network Classeasy Compose Multiple Transformsmedium

Concept primer

1. Background Knowledge

Tensor Shapes in PyTorch: Neural networks process multi-dimensional tensors. CNN feature maps typically have shape (N, C, H, W) where:

$N$ : batch size (number of samples)
$C$ : channels (feature maps)
$H, W$ : height, width (spatial dimensions)

nn.Flatten Purpose: Converts (N, C, H, W) to (N, C·H·W) for fully connected layers. Preserves batch dimension (dim=0) while collapsing all others into a feature vector of size $C \times H \times W$ .

Mathematical Transformation:

Input: (N, C, H, W) → Total elements per sample = C × H × W
Output: (N, C×H×W)

torch.randn: Generates tensors with i.i.d. samples from $\mathcal{N}(0,1)$ , standard for testing layer behavior.

2. Algorithm Approach

Core Pattern: Create → Transform → Analyze shapes

1. tensor = torch.randn(batch_size, channels, height, width)
2. flatten_layer = nn.Flatten()
3. flattened = flatten_layer(tensor)
4. Compute: input_shape = list(tensor.shape), output_shape = list(flattened.shape)
5. flattened_features = channels * height * width

nn.Flatten Parameters (defaults work here):

start_dim=1: Begin flattening from dimension 1 (preserves dim=0)
end_dim=-1: Flatten to last dimension

3. Step-by-Step Strategy

import torch
import torch.nn as nn

def apply_flatten(batch_size, channels, height, width):
 # Step 1: Create input tensor
 input_tensor = torch.randn(batch_size, channels, height, width)

 # Step 2: Apply nn.Flatten()
 flatten_layer = nn.Flatten()
 output_tensor = flatten_layer(input_tensor)

 # Step 3: Extract shapes as lists
 input_shape = list(input_tensor.shape)
 output_shape = list(output_tensor.shape)

 # Step 4: Calculate flattened features (per sample)
 flattened_features = channels * height * width

 # Step 5: Return dictionary
 return {
 "input_shape": input_shape,
 "output_shape": output_shape,
 "flattened_features": flattened_features
 }

Verification:

apply_flatten(2, 3, 4, 4) → {'input_shape': [2, 3, 4, 4], 'output_shape': [2, 48], 'flattened_features': 48}
# Since 3 × 4 × 4 = 48 features per sample

4. Common Pitfalls

Mistake	Why Wrong	Fix
tensor.view(-1)	Flattens entire batch to 1D	Use nn.Flatten(start_dim=1)
list(tensor.shape) → tensor.shape	Returns tuple, not list	Convert with list(tensor.shape)
flattened_features = output_shape	Brittle (assumes exact shape)	Compute channels * height * width
Forgetting import torch.nn as nn	nn.Flatten() undefined	Add import
torch.rand() instead of torch.randn()	Uniform [0,1] vs. Gaussian	Use torch.randn() per constraints

Shape Debugging Tip: Always print tensor.shape before/after operations.

5. Time & Space Complexity

Time Complexity: $O(N \times C \times H \times W)$

torch.randn: $O(V)$ where $V = NCHW$ is total volume
nn.Flatten: $O(V)$ (memory copy/reshape, no computation)
Shape operations: $O(1)$

Space Complexity: $O(N \times C \times H \times W)$

Input tensor: O(NCHW)
Output tensor: O(N × (CHW)) = O(NCHW) [same size]
Dictionary: O(1)
Total: O(NCHW)

Key Insight: Flattening preserves total elements ( $N \times C \times H \times W$ ) but changes layout from 4D to 2D for linear layer compatibility.

This covers all prerequisites for solving CNN→FC transition problems in PyTorch!

solution.py

Python 3.13

Test results