models/VGG16Classifier.py

import os
from torchvision.models import vgg16
import torch.nn.init as init
import torch.nn as nn
import torch


class VGG16Classifier(nn.Module):
    def __init__(self, num_classes: int = 14):
        super(VGG16Classifier, self).__init__()
        self.vgg16 = vgg16(pretrained=False)

        # Replace the classifier layer
        num_features = self.vgg16.classifier[6].in_features
        self.vgg16.classifier[6] = nn.Linear(num_features, num_classes)

        # 初始化权重
        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.Linear):
                    init.xavier_normal_(m.weight)
                    if m.bias is not None:
                        init.zeros_(m.bias)

    def forward(self, x):
        x = self.vgg16(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", "VGG16Classifier.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