| import torch |
| import torchvision |
| from torchvision.models import EfficientNet_B2_Weights |
| from torch import nn |
|
|
| def create_model(num_classes=7): |
| weights = EfficientNet_B2_Weights.DEFAULT |
| model = torchvision.models.efficientnet_b2(weights=weights) |
|
|
| for param in model.parameters(): |
| param.requires_grad = False |
|
|
| model.classifier = nn.Sequential( |
| nn.Dropout(p=0.3), |
| nn.Linear(model.classifier[1].in_features, num_classes) |
| ) |
|
|
| return model |
|
|
| def load_model(weights_path="model/effnetb2_dermamnist.pth"): |
| model = create_model(num_classes=7) |
| model.load_state_dict(torch.load(weights_path, map_location=torch.device("cpu"))) |
| model.eval() |
| return model |
|
|
