from torch import manual_seed, nn from torchvision import transforms, models def create_model_alexnet(num_classes:int=2, seed:int=42): """Creates model and transforms. Args: num_classes (int, optional): number of classes in the classifier head, defaults to 2. seed (int, optional): random seed value. Defaults to 42. Returns: model (torch.nn.Module): Alexnet model. transforms (torchvision.transforms): Alexnet image transforms. """ # Create Alexnet pretrained weights, transforms and model weights = models.AlexNet_Weights.IMAGENET1K_V1.DEFAULT auto_transform = weights.transforms() model_alexnet = models.alexnet(weights=weights) # Freeze all layers in base model for param in model_alexnet.parameters(): param.requires_grad = False # Change classifier head with random seed for reproducibility manual_seed(seed) model_alexnet.classifier[6] = nn.Linear(4096, out_features=num_classes) return model_alexnet, auto_transform