Datatrooper/posters_classification

1from torchvision import models as models2import torch.nn as nn3def model(pretrained, requires_grad):4    model = models.resnet50(progress=True, pretrained=pretrained)5    # to freeze the hidden layers6    if requires_grad == False:7        for param in model.parameters():8            param.requires_grad = False9    # to train the hidden layers10    elif requires_grad == True:11        for param in model.parameters():12            param.requires_grad = True13    # make the classification layer learnable14    # we have 25 classes in total15    model.fc = nn.Linear(2048, 25)16    return model