Forward pass computation in PyTorch - ML Experiment: Train & Evaluate

Experiment - Forward pass computation

Problem:You have a simple neural network model in PyTorch that classifies images into 10 classes. The model is defined, but you want to understand how the forward pass works and verify that the output shape and values are correct.

Current Metrics:No training done yet. The model outputs raw scores (logits) for each class. Output shape is (batch_size, 10).

Issue:You are unsure if the forward pass is implemented correctly and if the output tensor shape and values make sense.

Your Task

Run a forward pass on a batch of dummy input data and verify the output shape is (batch_size, 10). Also, check the output values to confirm they are raw scores (not probabilities).

Use PyTorch only.

Do not train the model; just run a forward pass.

Use a batch size of 4 and input size matching the model input.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

PyTorch

import torch
import torch.nn as nn

# Define a simple neural network
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear1 = nn.Linear(28*28, 128)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(128, 10)  # 10 classes

    def forward(self, x):
        x = self.flatten(x)
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)  # Output logits
        return x

# Instantiate the model
model = SimpleNet()

# Create dummy input: batch_size=4, channels=1, height=28, width=28
dummy_input = torch.randn(4, 1, 28, 28)

# Forward pass
output = model(dummy_input)

# Print output shape and values
print(f"Output shape: {output.shape}")
print(f"Output values (first sample): {output[0]}")

Defined a simple neural network with flatten, linear, and ReLU layers.

Created dummy input tensor with batch size 4 and correct input shape.

Ran forward pass using model(dummy_input).

Printed output shape and first sample output values to verify.

Results Interpretation

Before: No forward pass run, unsure about output shape and values.

After: Output shape confirmed as (4, 10), matching batch size and number of classes. Output values are raw scores (logits), not probabilities.

The forward pass transforms input data through the network layers to produce raw scores for each class. Checking output shape and values helps confirm the model is implemented correctly before training.

Bonus Experiment

Apply a softmax function to the output logits to convert them into probabilities and verify the probabilities sum to 1 for each sample.

💡 Hint

Use torch.nn.functional.softmax on the output tensor along the class dimension (dim=1). Then sum the probabilities for each sample to check they equal 1.