Spaces:
Build error
Build error
import torchvision | |
from torch import nn | |
def create_vit_b_16_swag(num_classes:int=101, seed:int=42): | |
# 1. Get the base mdoel with pretrained weights and send to target device | |
weights = torchvision.models.ViT_B_16_Weights.IMAGENET1K_SWAG_E2E_V1 | |
transforms = weights.transforms() | |
model = torchvision.models.vit_b_16(weights=weights)#.to(device) | |
# 2. Freeze the base model layers | |
for param in model.parameters(): | |
param.requires_grad = False | |
# 3. Change the heads | |
model.heads = nn.Linear(in_features=768, | |
out_features=101)#.to(device) | |
# 5. Give the model a name | |
model.name = "vit_b_16_swag" | |
print(f"[INFO] Created new {model.name} model.") | |
return model, transforms | |