Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| class ResNet(nn.Module): | |
| def __init__(self, in_channels: int, num_classes: int): | |
| """ResNet9""" | |
| super().__init__() | |
| self.conv1 = ConvBlock(in_channels, 64) | |
| self.conv2 = ConvBlock(64, 128, pool=True) | |
| self.res1 = nn.Sequential( | |
| ConvBlock(128, 128), | |
| ConvBlock(128, 128) | |
| ) | |
| self.conv3 = ConvBlock(128, 256) | |
| self.conv4 = ConvBlock(256, 512, pool=True) | |
| self.res2 = nn.Sequential( | |
| ConvBlock(512, 512), | |
| ConvBlock(512, 512) | |
| ) | |
| self.classifier = nn.Sequential( | |
| nn.MaxPool2d(kernel_size=(4, 4)), | |
| nn.AdaptiveAvgPool2d(1), | |
| nn.Flatten(), | |
| nn.Linear(512, 128), | |
| nn.Dropout(0.25), | |
| nn.Linear(128, num_classes), | |
| nn.Dropout(0.25), | |
| ) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = self.conv1(x) | |
| x = self.conv2(x) | |
| x = self.res1(x) + x #skip | |
| x = self.conv3(x) | |
| x = self.conv4(x) | |
| x = self.res2(x) + x #skip | |
| prediction = self.classifier(x) | |
| return prediction | |
| class ConvBlock(nn.Module): | |
| def __init__(self, in_channels: int, out_channels: int, pool: bool = False, pool_no: int = 2): | |
| super().__init__() | |
| self.in_channels = in_channels | |
| self.out_channels = out_channels | |
| self.pool = pool | |
| self.pool_no = pool_no | |
| if self.pool: | |
| self.pool_block = nn.Sequential( | |
| nn.ReLU(inplace=True), | |
| nn.MaxPool2d(self.pool_no) | |
| ) | |
| else: | |
| self.pool_block = nn.Sequential( | |
| nn.ReLU(inplace=True), | |
| ) | |
| self.block = nn.Sequential( | |
| nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), | |
| nn.BatchNorm2d(out_channels), | |
| self.pool_block | |
| ) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = self.block(x) | |
| return x |