Multi-class classification model

With a template for a binary classification model in place, you can now build on it to design a multi-class classification model. The model should handle different numbers of classes via a parameter, allowing you to tailor the model to a specific multi-class classification task in the future.

The packages torch and torch.nn as nn have been imported. All image sizes are 64x64 pixels.

This exercise is part of the course

Deep Learning for Images with PyTorch

View Course

Exercise instructions

Define the __init__ method including self and num_classes as parameters.
Create a fully connected layer with the input size of 16*32*32 and the number of classes num_classes as output.
Create an activation function softmax with dim=1.

Hands-on interactive exercise

Have a go at this exercise by completing this sample code.

class MultiClassImageClassifier(nn.Module):
  
    # Define the init method
    def ____(____, ____):
        super(MultiClassImageClassifier, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()

        # Create a fully connected layer
        self.fc = ____(____, ____)
        
        # Create an activation function
        self.softmax = ____(____)

Edit and Run Code