File size: 1,140 Bytes
18b0c85 5cc29b2 18b0c85 90dd800 18b0c85 5cc29b2 18b0c85 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | import torch
import torch.nn as nn
import numpy as np
import os
from torchvision import transforms
from torchvision.models import resnet18, ResNet18_Weights
class ResNet18_M3(nn.Module):
def __init__(self):
super().__init__()
self.model = resnet18(weights=ResNet18_Weights.DEFAULT)
self.model.fc = nn.Sequential(
nn.Linear(self.model.fc.in_features, 200),
nn.ReLU(),
nn.Dropout(p=0.3),
nn.Linear(200, 100),
nn.ReLU(),
nn.Dropout(p=0.3),
nn.Linear(100, 10)
)
def forward(self, x):
return self.model(x)
@staticmethod
def get_instance():
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
weights_path = os.path.join(os.path.dirname(__file__), "resnet18_m3_best.pth")
model = ResNet18_M3().to(DEVICE)
checkpoint = torch.load(
weights_path,
map_location=DEVICE,
weights_only=False
)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()
return model
|