| import torch | |
| import torchvision | |
| from torch import nn | |
| def create_effnet_b2_model(num_classes : int = 3, | |
| seed : int = 42): | |
| effnetb2_weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT | |
| effnetb2_transforms = effnetb2_weights.transforms() | |
| effnetb2 = torchvision.models.efficientnet_b2(weights=effnetb2_weights) | |
| for p in effnetb2.parameters(): | |
| p.requires_grad = False | |
| torch.manual_seed(seed) | |
| #torch.cuda.manual_seed(seed) | |
| effnetb2.classifier = nn.Sequential( | |
| torch.nn.Dropout(p=0.3, | |
| inplace=True), | |
| torch.nn.Linear(in_features=1408, | |
| out_features=num_classes, | |
| bias=True) | |
| ) | |
| return effnetb2, effnetb2_transforms | |