Spaces:
Sleeping
Sleeping
| import torch | |
| import torchvision | |
| from torch import nn | |
| def create_densenet161(output_shape: int = 4, | |
| seed: int = 42): | |
| """ | |
| Creates an DenseNet161 as a feature extractor setting all parameters as not trainable | |
| and changing the shape of the classifier layer equals to 'output_shape'. | |
| Also creates a corresponding 'torchivision.transforms' trasnformer to preprocess data | |
| Args: | |
| output_shape (int): number of classes | |
| seed (int): sedd for reproducibility | |
| Returns: | |
| model: DenseNet161 with 'output_shape' output neurons | |
| model_transformer: DenseNet161 transformer | |
| """ | |
| if seed: | |
| torch.manual_seed(seed) | |
| torch.cuda.manual_seed(seed) | |
| model = torchvision.models.densenet161(weights = torchvision.models.DenseNet161_Weights.DEFAULT) | |
| model_transformer = torchvision.models.DenseNet161_Weights.DEFAULT.transforms() | |
| for param in model.parameters(): | |
| param.requires_grad = False | |
| model.classifier = nn.Linear(in_features = 2208, out_features = output_shape, bias=True) | |
| return model, model_transformer | |