Reshaping (view, reshape, squeeze, unsqueeze) in PyTorch - ML Experiment: Train & Evaluate

Experiment - Reshaping (view, reshape, squeeze, unsqueeze)

Problem:You have a tensor representing a batch of images with shape (32, 3, 28, 28). You want to prepare this tensor for a fully connected neural network by flattening each image into a vector. Currently, the tensor shape is not compatible with the network input.

Current Metrics:Tensor shape before reshaping: torch.Size([32, 3, 28, 28])

Issue:The tensor shape is 4D, but the fully connected layer expects 2D input (batch_size, features). Without reshaping, the model will raise an error.

Your Task

Reshape the input tensor from shape (32, 3, 28, 28) to (32, 3*28*28) so it can be fed into a fully connected layer. Then, after the layer, reshape the output back to (32, 10, 1, 1) to simulate a feature map. Use PyTorch functions view, reshape, squeeze, and unsqueeze appropriately.

Do not change the batch size (32).

Use PyTorch tensor operations only.

Avoid using loops.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch

# Original tensor simulating a batch of images
x = torch.randn(32, 3, 28, 28)
print(f"Original shape: {x.shape}")

# Flatten each image to a vector of size 3*28*28
x_flat = x.view(32, -1)  # or x.reshape(32, -1)
print(f"After flattening: {x_flat.shape}")

# Simulate a fully connected layer output with 10 features
fc_output = torch.randn(32, 10)
print(f"Fully connected output shape: {fc_output.shape}")

# Reshape fc_output to (32, 10, 1, 1) to simulate feature maps
fc_reshaped = fc_output.unsqueeze(-1).unsqueeze(-1)
print(f"After unsqueeze: {fc_reshaped.shape}")

# Remove the last two dimensions if needed
fc_squeezed = fc_reshaped.squeeze(-1).squeeze(-1)
print(f"After squeeze: {fc_squeezed.shape}")

Used view to flatten the 4D tensor to 2D for fully connected layer input.

Used unsqueeze twice to add two singleton dimensions to the fully connected output.

Used squeeze twice to remove singleton dimensions from the tensor.

Results Interpretation

Before: Input tensor shape was (32, 3, 28, 28), which is 4D and incompatible with fully connected layers.

After: Flattened tensor shape is (32, 2352), suitable for fully connected layers. Then reshaped output to (32, 10, 1, 1) using unsqueeze, and back to (32, 10) using squeeze.

Reshaping tensors is essential to prepare data for different layers in neural networks. Using view or reshape changes tensor shape without copying data. Squeeze removes dimensions of size 1, and unsqueeze adds such dimensions. These operations help adapt tensor shapes flexibly.

Bonus Experiment

Try reshaping the tensor using only reshape instead of view, and verify if the output shapes remain the same.

💡 Hint

Both view and reshape can change tensor shapes, but reshape can handle non-contiguous tensors better. Check if the tensor is contiguous before using view.