1. Learn
    2. /
  2. Courses
    3. /
  3. Deep Learning for Images with PyTorch
Connected

Exercise

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.

Instructions

100 XP
  • 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.