|
|
import torch |
|
|
import torch.nn as nn |
|
|
|
|
|
from torchvision.models import densenet121, DenseNet121_Weights |
|
|
from torchvision.models import resnet50, ResNet50_Weights |
|
|
from torchvision.models import efficientnet_v2_m, EfficientNet_V2_M_Weights |
|
|
from torchvision.models import alexnet, AlexNet_Weights |
|
|
from torchvision.models import vgg16, VGG16_Weights |
|
|
from torchvision.models import vgg19, VGG19_Weights |
|
|
|
|
|
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') |
|
|
|
|
|
def changedClassifierLayer(model, modelName, N_CLASSES=10): |
|
|
for param in model.parameters(): |
|
|
param.requires_grad = False |
|
|
|
|
|
if modelName == "DenseNet121": |
|
|
num_input = model.classifier.in_features |
|
|
|
|
|
elif modelName == "ResNet50": |
|
|
num_input = model.fc.in_features |
|
|
|
|
|
elif modelName == "EfficientNet-V2-M" or modelName == "AlexNet": |
|
|
num_input = model.classifier[1].in_features |
|
|
|
|
|
elif modelName == "VGG19" or modelName == "VGG16": |
|
|
num_input = model.classifier[0].in_features |
|
|
|
|
|
classifier = nn.Sequential( |
|
|
nn.Linear(num_input, 256), |
|
|
nn.ReLU(), |
|
|
nn.Dropout(0.2), |
|
|
nn.Linear(256, 128), |
|
|
nn.ReLU(), |
|
|
nn.Dropout(0.2), |
|
|
nn.Linear(128, N_CLASSES), |
|
|
nn.LogSoftmax(dim=1) |
|
|
) |
|
|
|
|
|
if modelName == "ResNet50": |
|
|
model.fc = classifier |
|
|
else: |
|
|
model.classifier = classifier |
|
|
|
|
|
efficientnet_weights_path = 'models/EfficientNet-V2-M.pth' |
|
|
densenet_weights_path = 'models/DenseNet121.pth' |
|
|
resnet_weights_path = 'models/ResNet50.pth' |
|
|
alexnet_weights_path = 'models/AlexNet.pth' |
|
|
vgg16_weights_path = 'models/VGG16.pth' |
|
|
vgg19_weights_path = 'models/VGG19.pth' |
|
|
|
|
|
efficientnetV2M_model = efficientnet_v2_m(weights=EfficientNet_V2_M_Weights.IMAGENET1K_V1) |
|
|
densenet_model = densenet121(weights=DenseNet121_Weights.IMAGENET1K_V1) |
|
|
resnet_model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2) |
|
|
alexnet_model = alexnet(weights=AlexNet_Weights.IMAGENET1K_V1) |
|
|
vgg16_model = alexnet(weights=VGG16_Weights.IMAGENET1K_V1) |
|
|
vgg19_model = alexnet(weights=VGG19_Weights.IMAGENET1K_V1) |
|
|
|
|
|
changedClassifierLayer(efficientnetV2M_model, "EfficientNet-V2-M") |
|
|
changedClassifierLayer(densenet_model, "DenseNet121") |
|
|
changedClassifierLayer(resnet_model, "ResNet50") |
|
|
changedClassifierLayer(alexnet_model, "AlexNet") |
|
|
changedClassifierLayer(vgg16_model, "VGG16") |
|
|
changedClassifierLayer(vgg19_model, "VGG19") |
|
|
|
|
|
efficientnetV2M_model.load_state_dict(torch.load(efficientnet_weights_path)) |
|
|
densenet_model.load_state_dict(torch.load(densenet_weights_path)) |
|
|
resnet_model.load_state_dict(torch.load(resnet_weights_path)) |
|
|
alexnet_model.load_state_dict(torch.load(alexnet_weights_path)) |
|
|
vgg16_model.load_state_dict(torch.load(vgg16_weights_path)) |
|
|
vgg19_model.load_state_dict(torch.load(vgg19_weights_path)) |
|
|
|
|
|
class EnsembleModel(nn.Module): |
|
|
def __init__(self, model_list, weights=None): |
|
|
super(EnsembleModel, self).__init__() |
|
|
self.models = nn.ModuleList(model_list) |
|
|
self.weights = weights |
|
|
|
|
|
def forward(self, x): |
|
|
outputs = [model(x.to(next(model.parameters()).device)) for model in self.models] |
|
|
|
|
|
if self.weights is None: |
|
|
ensemble_output = torch.mean(torch.stack(outputs), dim=0) |
|
|
|
|
|
|
|
|
else: |
|
|
weighted_outputs = torch.stack([w * output for w, output in zip(self.weights, outputs)]) |
|
|
ensemble_output = torch.sum(weighted_outputs, dim=0) |
|
|
|
|
|
return ensemble_output |
|
|
|
|
|
models_list = [ |
|
|
efficientnetV2M_model, |
|
|
densenet_model, |
|
|
resnet_model, |
|
|
alexnet_model, |
|
|
vgg16_model, |
|
|
vgg19_model |
|
|
] |
|
|
|
|
|
ensemble_model = EnsembleModel(models_list) |
|
|
|
|
|
|
