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 |