File size: 1,904 Bytes
9f20d50 |
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 |
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 |