Team Ai
Apppublic

JohnPinto/Human_Activity_Recognition-HAR-Video_Classification-HMDB51-Dataset

sourceHugging Facemitupdated 4y agoView on Hugging Face
2likes
model.py32 linesDownload Raw Back to root
1import torch2import torchvision3 4def create_model(num_classes: int, seed: int = 42):5    """6    A function to create a model.7    Parameters:8        num_classes: int, A integer for toal number of classes.9        seed: int(default: 42), A random seed value.10    Returns: 11        model: A feature extracted model for video classification.12        transforms: A torchvision transform is returned which was used in the pretrained model.    13    """14    # Creating model, weights and transforms15    weights = torchvision.models.video.MViT_V2_S_Weights.DEFAULT16    transforms = weights.transforms()17    model = torchvision.models.video.mvit_v2_s(weights=weights)18    19    # Freezing the model layers20    for params in model.parameters():21        params.requires_grad = False22        23    # Changing the fully Conncected head layer24    torch.manual_seed(seed)25    dropout_layer = model.head[0]26    in_features = model.head[1].in_features27    model.head = torch.nn.Sequential(28        dropout_layer,29        torch.nn.Linear(in_features=in_features, out_features=num_classes, bias=True)30    )31    return model, transforms32