A logit is a raw, unnormalized real-valued score assigned to each class before scores become probabilities through softmax. A classifier therefore returns one logit per possible label rather than a probability distribution. section 3.3.5 section 3.4.3
A vision-transformer classifier produces logits through this pipeline:
- Pass images through embeddings that combine token embeddings with positional encodings. section 3.3.5
- Run the embeddings through the encoder-only vision transformer. section 3.3.5
- Read the representation associated with the CLS token, which aggregates information across the image. section 3.3.5
- Apply a linear classifier head to produce one logit for each of the 10 CIFAR-10 classes. section 3.3.5
- Select the class with the highest logit for a prediction, or apply softmax when probability distributions are needed. section 3.4.3 section 3.3.5
Classifying a test batch
A trained ViT receives 32 CIFAR-10 test images and returns a logit tensor. Applying argmax across the class dimension produces 32 numerical predictions, which can be mapped to labels such as ship, plane, frog, or car. section 3.4.3
Listing 3.7 Building a classifier based on the ViT
class ViTForClassfication(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.image_size = config.image_size
self.hidden_size = config.hidden_size
self.num_classes = config.num_classes
self.embedding = Embeddings(config)
self.encoder = Encoder(config)
self.classifier = nn.Linear(self.hidden_size,
self.num_classes) #1
self.apply(self._init_weights)
def forward(self, x, output_attentions=False):
embedding_output = self.embedding(x) #2
encoder_output, all_attentions = self.encoder(\ #3
embedding_output,output_attentions=output_attentions)
logits = self.classifier(encoder_output[:, 0, :]) #4
if not output_attentions:
return (logits, None)
else:
return (logits, all_attentions)
Inspect the classifier implementation to see the linear head transform the CLS-token representation into class logits. section 3.3.5
Listing 3.10 Using the trained ViT to classify images
import math
import torch.nn.functional as F
model.load_state_dict(torch.load('files/ViT.pth'))
model.eval()
with torch.no_grad():
batch=next(iter(testloader)) #1
batch = [t.to(device) for t in batch]
images, labels = batch
logits, attention_maps = model(images,
output_attentions=True) #2
predictions = torch.argmax(logits, dim=1) #3
print(predictions)
print([names[i] for i in predictions.tolist()]) #4
Follow the inference code that loads trained weights, obtains logits, and selects predictions with argmax. section 3.4.3
