Upload Models.py

Models.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+import timm
+
+# ResNet50 Model
+class ResNet(nn.Module):
+    def __init__(self, num_classes, is_freeze=True):
+        super(ResNet, self).__init__()
+
+        self.num_classes = num_classes
+        self.is_freeze = is_freeze
+        self.base_model = timm.create_model('resnet50', pretrained=True)
+
+        if self.is_freeze:
+            for param in self.base_model.parameters():
+                param.requires_grad = False
+      
+        self.base_model.fc = nn.Linear(2048, self.num_classes)
+
+    def forward(self, x):
+        x = self.base_model(x)
+        return x
+
+# EfficientNet Model
+class EfficientNet(nn.Module):
+    def __init__(self, num_classes):
+        super(EfficientNet, self).__init__()
+
+        self.num_classes = num_classes
+        self.base_model = timm.create_model('efficientnet_b0', pretrained=True)
+        self.base_model.classifier = nn.Linear(1280, self.num_classes)
+
+    def forward(self, x):
+        x = self.base_model(x)
+        return x
+
+# BaseLine Model
+class BaseLine(nn.Module):
+    def __init__(self, num_classes):
+        super(BaseLine, self).__init__()
+
+        self.Conv1 = nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=1)
+        self.Conv2 = nn.Conv2d(96, 256, kernel_size=5, padding=2)
+        self.Conv3 = nn.Conv2d(256, 384, kernel_size=3, padding=1)
+        self.Conv4 = nn.Conv2d(384, 256, kernel_size=3, padding=1)
+
+        self.Linear1 = nn.Linear(2304, 512)
+        self.Linear3 = nn.Linear(512, num_classes)
+
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(p=0.5)
+
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2)
+        self.flatten = nn.Flatten()
+
+    def forward(self, x):
+        x = self.Conv1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.Conv2(x)
+        x = self.maxpool(x)
+
+        x = self.Conv3(x)
+        x = self.Conv4(x)
+        x = self.maxpool(x)
+
+        x = self.flatten(x)
+        x = self.Linear1(x)
+        x = self.relu(x)
+        x = self.dropout(x)
+
+        x = self.Linear3(x)
+        return x