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