Vrk/MultiClass-Classifier

1import torch2import torch.nn as nn3import torch.optim as optim4import torch.nn.functional as F5import timm6 7# ResNet50 Model8class ResNet(nn.Module):9    def __init__(self, num_classes, is_freeze=True):10        super(ResNet, self).__init__()11 12        self.num_classes = num_classes13        self.is_freeze = is_freeze14        self.base_model = timm.create_model('resnet50', pretrained=True)15 16        if self.is_freeze:17            for param in self.base_model.parameters():18                param.requires_grad = False19      20        self.base_model.fc = nn.Linear(2048, self.num_classes)21 22    def forward(self, x):23        x = self.base_model(x)24        return x25 26# EfficientNet Model27class EfficientNet(nn.Module):28    def __init__(self, num_classes):29        super(EfficientNet, self).__init__()30 31        self.num_classes = num_classes32        self.base_model = timm.create_model('efficientnet_b0', pretrained=True)33        self.base_model.classifier = nn.Linear(1280, self.num_classes)34 35    def forward(self, x):36        x = self.base_model(x)37        return x38 39# BaseLine Model40class BaseLine(nn.Module):41    def __init__(self, num_classes):42        super(BaseLine, self).__init__()43 44        self.Conv1 = nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=1)45        self.Conv2 = nn.Conv2d(96, 256, kernel_size=5, padding=2)46        self.Conv3 = nn.Conv2d(256, 384, kernel_size=3, padding=1)47        self.Conv4 = nn.Conv2d(384, 256, kernel_size=3, padding=1)48 49        self.Linear1 = nn.Linear(2304, 512)50        self.Linear3 = nn.Linear(512, num_classes)51 52        self.relu = nn.ReLU()53        self.dropout = nn.Dropout(p=0.5)54 55        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2)56        self.flatten = nn.Flatten()57 58    def forward(self, x):59        x = self.Conv1(x)60        x = self.relu(x)61        x = self.maxpool(x)62 63        x = self.Conv2(x)64        x = self.maxpool(x)65 66        x = self.Conv3(x)67        x = self.Conv4(x)68        x = self.maxpool(x)69 70        x = self.flatten(x)71        x = self.Linear1(x)72        x = self.relu(x)73        x = self.dropout(x)74 75        x = self.Linear3(x)76        return x