Add softmax layer to vision transformer model

vision
1

I am trying to add a softmax layer to a vit-b model with 10 outputs.
While it works calling torch.softmax directly on the model output, I do not get softmax output when adding the softmax as a fully connected layer:

import torchvision
import torch

torch.manual_seed(1)
image = torch.rand(1,3,224,224)
model = torchvision.models.vit_b_16(weights="DEFAULT")
model.heads.head = torch.nn.Linear(in_features=768, out_features=10, bias=True)
model.fc = torch.nn.Softmax(10)
model(image)

This does not seem to apply softmax, the output is

tensor([[ 0.0204, -0.0550, -0.7651, -0.2418,  0.0715,  0.1525, -0.1907, -0.1199,
          0.0427, -0.2585]], grad_fn=<AddmmBackward0>)

What am I doing wrong?

2

vit_b_16 doesn’t use a .fc module so you are creating a new attribute in:

model.fc = torch.nn.Softmax(10)

which is never used.
You could replace the model.heads.head with an nn.Sequential container containing the new nn.Linear as well as the nn.Softmax layer.

3

Indeed, this worked.

model.heads.head = torch.nn.Sequential(torch.nn.Linear(in_features=768, out_features=10, bias=True), torch.nn.Softmax(dim=1))

Thank you very much for your help!