unit 4.4 - Fixing our first example

Do you remember the neural network we asked Gemini to create for us in lecture 0.0? It is below!

We now know enough to be able to fix it and make it run!

This is your homework!

import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.nn.functional as F

class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()

        # Convolutional layers
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)

        # Pooling layers
        self.pool1 = nn.AvgPool2d(2)
        self.pool2 = nn.AvgPool2d(2)

        # Fully connected layers
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # Convolutional layers
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        x = F.relu(self.conv2(x))
        x = self.pool2(x)

        # Fully connected layers
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)

        return x

# Load the CIFAR-10 dataset
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transforms.ToTensor())
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transforms.ToTensor())

# Create the LeNet-5 model
model = LeNet5()

# Define the loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# Train the model
for epoch in range(10):
    # Train the model on the training dataset
    for i, (data, labels) in enumerate(train_dataset):
        # Forward pass
        outputs = model(data)

        # Compute the loss
        loss = criterion(outputs, labels)

        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Print the progress
        if (i + 1) % 100 == 0:
            print('Epoch [{}/{}], Batch [{}/{}], Loss: {:.6f}'.format(epoch + 1, 10, i + 1, len(train_dataset), loss.item()))

# Evaluate the model on the test dataset
correct = 0
total = 0
with torch.no_grad():
    for data, labels in test_dataset:
        outputs = model(data)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the network on the test images: {} %'.format(100 * correct / total))

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[1], line 58
     55 outputs = model(data)
     57 # Compute the loss
---> 58 loss = criterion(outputs, labels)
     60 # Backward pass
     61 optimizer.zero_grad()

File /opt/homebrew/lib/python3.13/site-packages/torch/nn/modules/module.py:1773, in Module._wrapped_call_impl(self, *args, **kwargs)
   1771     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1772 else:
-> 1773     return self._call_impl(*args, **kwargs)

File /opt/homebrew/lib/python3.13/site-packages/torch/nn/modules/module.py:1784, in Module._call_impl(self, *args, **kwargs)
   1779 # If we don't have any hooks, we want to skip the rest of the logic in
   1780 # this function, and just call forward.
   1781 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1782         or _global_backward_pre_hooks or _global_backward_hooks
   1783         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1784     return forward_call(*args, **kwargs)
   1786 result = None
   1787 called_always_called_hooks = set()

File /opt/homebrew/lib/python3.13/site-packages/torch/nn/modules/loss.py:1310, in CrossEntropyLoss.forward(self, input, target)
   1309 def forward(self, input: Tensor, target: Tensor) -> Tensor:
-> 1310     return F.cross_entropy(
   1311         input,
   1312         target,
   1313         weight=self.weight,
   1314         ignore_index=self.ignore_index,
   1315         reduction=self.reduction,
   1316         label_smoothing=self.label_smoothing,
   1317     )

File /opt/homebrew/lib/python3.13/site-packages/torch/nn/functional.py:3462, in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)
   3460 if size_average is not None or reduce is not None:
   3461     reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 3462 return torch._C._nn.cross_entropy_loss(
   3463     input,
   3464     target,
   3465     weight,
   3466     _Reduction.get_enum(reduction),
   3467     ignore_index,
   3468     label_smoothing,
   3469 )

TypeError: cross_entropy_loss(): argument 'target' (position 2) must be Tensor, not int