File size: 4,422 Bytes
60b0ddc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
import torch
import torch.nn as nn
import torchvision.models as models
# Improved Net
class Net(nn.Module):
def __init__(self, n_classes: int) -> None:
super(Net, self).__init__()
self.cnn_layers = nn.Sequential(
nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.AvgPool2d(kernel_size=2),
nn.Dropout(p=0.5),
nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
nn.AvgPool2d(kernel_size=2),
nn.Dropout(p=0.25),
nn.Conv2d(32, 16, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(16),
nn.ReLU(inplace=True),
nn.AvgPool2d(kernel_size=2),
nn.Dropout(p=0.125),
nn.Conv2d(16, 8, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(8),
nn.ReLU(inplace=True),
nn.AvgPool2d(kernel_size=2),
nn.Dropout(p=0.1),
nn.Conv2d(8, 4, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(4),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2),
nn.Dropout(p=0.05),
# New layer
nn.Conv2d(4, 4, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(4),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2),
nn.Dropout(p=0.05),
)
self.linear_layers = nn.Sequential(
nn.Linear(16, 256),
nn.Linear(256, n_classes)
)
# Defining the forward pass
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.cnn_layers(x)
# After our convolutional layers which are 2D, we need to flatten our
# input to be 1 dimensional, as the linear layers require this.
x = x.view(x.size(0), -1)
x = self.linear_layers(x)
return x
# Implementing Pre-Trained ResNet as a class
class ResNetModel(nn.Module):
def __init__(self, n_classes: int, pretrained: bool = True):
# Loading a pre-trained ResNet model
super(ResNetModel, self).__init__()
self.resnet = models.resnet34(pretrained=pretrained)
self.resnet.conv1 = nn.Conv2d(1, 64, kernel_size=4, stride=(2, 2), padding=(3, 3), bias=False)
num_ftrs = self.resnet.fc.in_features
self.resnet.fc = nn.Linear(num_ftrs, n_classes)
def forward(self, x):
# Forward pass through the ResNet model
return self.resnet(x)
# Implementing the Efficient Net Model with efficientnet_b0
class EfficientNetModel(nn.Module):
def __init__(self, n_classes: int, version: str = 'b0', pretrained: bool = True):
super(EfficientNetModel, self).__init__()
# Loading a pretrained EfficientNet model
self.efficientnet = models.efficientnet_b0(pretrained=pretrained) if version == 'b0' else models.__dict__[f'efficientnet_{version}'](pretrained=pretrained)
# Adjusting the classifier to match the number of classes
num_ftrs = self.efficientnet.classifier[1].in_features
self.efficientnet.classifier[1] = nn.Linear(num_ftrs, n_classes)
def forward(self, x):
# Forward pass through the EfficientNet model
# Replicating the grayscale channel to have 3 channels
x = x.repeat(1, 3, 1, 1)
return self.efficientnet(x)
# Implementing the Efficient Net Model with efficientnet_b7
# NOTE: This model takes a lot of time to run
class EfficientNetModel_b7(nn.Module):
def __init__(self, n_classes: int, version: str = 'b0', pretrained: bool = True):
super(EfficientNetModel_b7, self).__init__()
# Loading a pretrained EfficientNet model
self.efficientnet = models.efficientnet_b7(pretrained=pretrained) if version == 'b7' else models.__dict__[f'efficientnet_{version}'](pretrained=pretrained)
# Adjusting the classifier to match the number of classes
num_ftrs = self.efficientnet.classifier[1].in_features
self.efficientnet.classifier[1] = nn.Linear(num_ftrs, n_classes)
def forward(self, x):
# Forward pass through the EfficientNet model
# Replicating the grayscale channel to have 3 channels
x = x.repeat(1, 3, 1, 1)
return self.efficientnet(x)
|