|
|
import os |
|
|
from torchvision.models.inception import inception_v3 |
|
|
import torch.nn.init as init |
|
|
import torch.nn as nn |
|
|
import torch |
|
|
|
|
|
|
|
|
class InceptionV3Classifier(nn.Module): |
|
|
def __init__(self, num_classes: int = 14): |
|
|
super(InceptionV3Classifier, self).__init__() |
|
|
self.inception = inception_v3( |
|
|
pretrained=False, num_classes=num_classes, aux_logits=False |
|
|
) |
|
|
|
|
|
|
|
|
if not self._load(): |
|
|
for m in self.modules(): |
|
|
if isinstance(m, nn.Conv2d): |
|
|
init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") |
|
|
if m.bias is not None: |
|
|
init.zeros_(m.bias) |
|
|
elif isinstance(m, nn.BatchNorm2d): |
|
|
init.ones_(m.weight) |
|
|
init.zeros_(m.bias) |
|
|
elif isinstance(m, nn.Linear): |
|
|
init.xavier_normal_(m.weight) |
|
|
if m.bias is not None: |
|
|
init.zeros_(m.bias) |
|
|
|
|
|
def forward(self, x): |
|
|
x = self.inception(x) |
|
|
return x |
|
|
|
|
|
def _load(self, filename: str = None) -> bool: |
|
|
if filename is None: |
|
|
current_work_dir = os.path.dirname(__file__) |
|
|
filename = os.path.join( |
|
|
current_work_dir, "best_pth", "InceptionV3Classifier.pth" |
|
|
) |
|
|
if not os.path.exists(filename): |
|
|
print("Model file does not exist.") |
|
|
return False |
|
|
self.load_state_dict(torch.load(filename)) |
|
|
print("Model loaded successfully.") |
|
|
return True |
|
|
|
