Use nn.Flatten Layer
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:
batch_size=2, channels=3, height=4, width=4
{'input_shape': [2, 3, 4, 4], 'output_shape': [2, 48], 'flattened_features': 48}344 = 48 features, batch size 2 preserved
Constraints:
- Use nn.Flatten()
- Batch dimension should be preserved
- Use torch.randn for random tensor
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×H×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 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×C×H×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×C×H×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×C×H×W) but changes layout from 4D to 2D for linear layer compatibility.
This covers all prerequisites for solving CNN→FC transition problems in PyTorch!