model.py

import torchvision
from torch import nn


def create_swin(class_count, version="tiny", use_v2=True, device="cpu"):
    v_string = "v2_" if use_v2 else ""
    weights = torchvision.models.get_weight(f"Swin_{v_string.upper()}{version[0].upper()}_Weights.DEFAULT")
    model = torchvision.models.get_model(f"swin_{v_string.lower()}{version[0].lower()}", weights=weights).to(device)
    input_features = model.head.in_features
    for param in model.parameters():
        param.requires_grad = False
    model.head = nn.Linear(input_features, class_count)
    return model, weights.transforms()