| from torch import nn | |
| import torch | |
| class resnet (nn.Module) : | |
| def __init__ (self, n_in,n_out,strid=1,padding=1) : | |
| super(resnet, self).__init__() | |
| self.cnn1 = nn.Conv2d(in_channels=n_in, out_channels=n_out, kernel_size=3, stride=strid, padding=1) | |
| self.cnn2 = nn.Conv2d(in_channels=n_out, out_channels=n_out, kernel_size=3, stride=1, padding=1) | |
| self.normal = nn.BatchNorm2d(n_out) | |
| self.normal2 = nn.BatchNorm2d(n_out) | |
| self.relu = nn.ReLU() | |
| self.relu2 = nn.ReLU() | |
| self.relu3 = nn.ReLU() | |
| self.shorcut = nn.Identity() | |
| if strid != 1 or n_in != n_out : | |
| self.shorcut = nn.Sequential( | |
| nn.Conv2d(in_channels=n_in, out_channels=n_out, kernel_size=1, stride=strid), | |
| nn.BatchNorm2d(n_out) | |
| ) | |
| def forward(self, x) : | |
| identity = self.shorcut(x) | |
| x = self.normal(self.cnn1(x)) | |
| x = self.relu(x) | |
| x = self.normal2(self.cnn2(x)) | |
| x = self.relu2(x) | |
| x = x + identity | |
| x = self.relu3(x) | |
| return x | |
| class Resnet (nn.Module) : | |
| def __init__ (self,num_class= 3) : | |
| super(Resnet, self).__init__() | |
| self.cnn1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=7, stride=2, padding=3) | |
| self.normal1 = nn.BatchNorm2d(64) | |
| self.relu = nn.ReLU() | |
| self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) | |
| self.resnet1 = resnet(64,64) | |
| self.resnet2 = resnet(64,128,2) | |
| self.resnet3 = resnet(128,256,2) | |
| self.resnet4 = resnet(256,512,2) | |
| self.pooling = nn.AdaptiveAvgPool2d((1,1)) | |
| self.linear1 = nn.Linear(512, num_class) | |
| def forward(self, x) : | |
| x = self.cnn1(x) | |
| x = self.normal1(x) | |
| x = self.relu(x) | |
| x = self.maxpool(x) | |
| x = self.resnet1(x) | |
| x = self.resnet2(x) | |
| x = self.resnet3(x) | |
| x = self.resnet4(x) | |
| x = self.pooling(x) | |
| x = torch.flatten(x,1) | |
| x = self.linear1(x) | |
| return x |