Model Configuration, Training, and Predictions for Classification

Get to know the model configuration and training implementation changes for classification problems along with the steps of making predictions.

Model configuration

In the chapter, Going Classy, we ended up with a lean model configuration part. We only need to define a model, an appropriate loss function, and an optimizer. Let us define a model that produces logits and uses BCEWithLogitsLoss as the loss function. Since we have two features, and we are producing logits instead of probabilities, our model has one layer and one layer alone: Linear(2, 1). We will keep using the SGD optimizer with a learning rate of 0.1 for now.

This is what the model configuration looks like for our classification problem:

Get hands-on with 1200+ tech skills courses.