Spaces:
Runtime error
Runtime error
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 |