File size: 704 Bytes
45d2464 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
import torch
import torchvision
from torch import nn
def create_vit_model(num_classes: int = 3,
seed: int = 42):
# Create ViT_B_16 pre-trained weights, transforms and model
weights = torchvision.models.ViT_B_16_Weights.DEFAULT
transforms = weights.transforms()
model = torchvision.models.vit_b_16(weights = weights)
# Freeze all of the base layers
for param in model.parameters():
param.requires_grad = False
# Change classifier head to suit our needs
torch.manual_seed(seed)
model.heads = nn.Sequential(nn.Linear(in_features = 768,
out_features = num_classes))
return model, transforms
|