Ordering of softmax in Convolutional Neural Network (CNN)

[UrDailyCs]

Hey daniel, i wanna ask if i were to use softmax in CNN multi-class classification problem . is this the right implementation?
assume that in_features and out_features is right, i wanna ask is my softmax placement correct? or i should use another softmax in the end?

  self.classifier = nn.Sequential(

        nn.Flatten(),
        nn.Linear(in_features = 32*8*8, out_features = 256),
        nn.Softmax(dim=1),
        nn.Linear(in_features = 256, out_features = output_shape)
)

[mrdbourke]

Hi @UrDailyCs,

Typically in PyTorch you'd use nn.Softmax() outside of the model (e.g. the model outputs the raw logits from nn.Linear() and then the softmax function gets used on them.

For example:

import torch
from torch import nn

# Create model
model =  nn.Sequential(
        nn.Flatten(),
        nn.Linear(in_features = 32*8*8, out_features = 256),
        nn.Linear(in_features = 256, out_features = output_shape)
)

# Put data through model
outputs = model(x)

# Use softmax on model outputs
pred_probs = torch.softmax(outputs)

That is the more standard approach.

However, you could try your approach and see what happens.

A good experiment would be to try both and compare the results.