Spaces:
Sleeping
Sleeping
| import torch.nn as nn | |
| import torchinfo | |
| import torch.nn.functional as F | |
| class convLayer(nn.Module): | |
| def __init__(self, l_input_c, | |
| l_output_c, bias=False, | |
| padding=1, stride=1, | |
| max_pooling=False, | |
| dropout=0): | |
| super (convLayer, self).__init__() | |
| self.convLayer = nn.Conv2d(in_channels=l_input_c, | |
| out_channels=l_output_c, | |
| kernel_size=(3, 3), | |
| stride=stride, | |
| padding= padding, | |
| padding_mode='replicate', | |
| bias=bias) | |
| self.max_pooling = None | |
| if(max_pooling == True): | |
| self.max_pooling = nn.MaxPool2d(2, 2) | |
| self.normLayer = nn.BatchNorm2d(l_output_c) | |
| self.activationLayer = nn.ReLU() | |
| self.dropout = None | |
| if(dropout > 0): | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x): | |
| x = self.convLayer(x) | |
| if (self.max_pooling is not None): | |
| x = self.max_pooling(x) | |
| x = self.normLayer(x) | |
| x = self.activationLayer(x) | |
| if (self.dropout is not None): | |
| x = self.dropout(x) | |
| return x | |
| class custBlock(nn.Module): | |
| def __init__(self, l_input_c, | |
| l_output_c, bias=False, | |
| padding=1, stride=1, | |
| max_pooling=True, | |
| dropout=0, residual_links=2): | |
| super (custBlock, self).__init__() | |
| self.conv_pool_block = convLayer(l_input_c=l_input_c, | |
| l_output_c=l_output_c, | |
| bias=bias, padding=padding, | |
| stride=stride, max_pooling=max_pooling, | |
| dropout=dropout) | |
| self.residual_block = None | |
| if(residual_links > 0): | |
| res_layer_seq = [] | |
| for link in range(0, residual_links): | |
| res_layer_seq.append( | |
| convLayer(l_input_c=l_output_c, | |
| l_output_c=l_output_c, | |
| bias=bias, padding=padding, | |
| stride=stride, max_pooling=False, | |
| dropout=dropout) | |
| ) | |
| self.residual_block = nn.Sequential(*res_layer_seq) | |
| def forward(self, x): | |
| x = self.conv_pool_block(x) | |
| if (self.residual_block is not None): | |
| tmp_x = x | |
| x = self.residual_block(x) | |
| x = x + tmp_x | |
| return x | |
| class custResNet(nn.Module): | |
| def __init__(self, dropout=0): | |
| super(custResNet, self).__init__() | |
| ##### Prep Block ##### | |
| # This block has a 3x3 convolution layer with stride=1, | |
| # padding=1 followed by batch normalization and RELU | |
| # 64 kernels are used in this block | |
| # By default, dropout is set to 0 | |
| self.prep_block = custBlock(l_input_c=3, l_output_c=64, | |
| max_pooling=False, dropout= dropout, | |
| residual_links=0 | |
| ) # output_size = 32, rf_out = 3 | |
| ##### Convolution Block - 1 ##### | |
| # This block in the first step has a 3x3 convolution layer with | |
| # stride=1, padding=1 followed by Max pooling, batch normalization | |
| # and ReLU. In the second step a network with residual links, with | |
| # each link having 3x3 convolution with stride=1, padding=1 | |
| # followed by batch normalization and ReLU. And in the third step, | |
| # the result from the first step and the result of the residual network | |
| # from the second step are added to make a skip connection. | |
| # 128 kernels are used in this block | |
| # By default, dropout is set to 0 | |
| self.block1 = custBlock(l_input_c=64, l_output_c=128, | |
| max_pooling=True, dropout= dropout, | |
| residual_links=2 | |
| ) # output_size = 16, rf_out = 13 | |
| ##### Convolution Block - 2 ##### | |
| # This block in the first step has a 3x3 convolution layer with | |
| # stride=1, padding=1 followed by Max pooling, batch normalization | |
| # and ReLU. | |
| # 256 kernels are used in this block | |
| # By default, dropout is set to 0 | |
| self.block2 = custBlock(l_input_c=128, l_output_c=256, | |
| max_pooling=True, dropout= dropout, | |
| residual_links=0 | |
| ) # output_size = 8, rf_out = 17 | |
| ##### Convolution Block - 3 ##### | |
| # This block in the first step has a 3x3 convolution layer with | |
| # stride=1, padding=1 followed by Max pooling, batch normalization | |
| # and ReLU. In the second step a network with residual links, with | |
| # each link having 3x3 convolution with stride=1, padding=1 | |
| # followed by batch normalization and ReLU. And in the third step, | |
| # the result from the first step and the result of the residual network | |
| # from the second step are added to make a skip connection. | |
| # 512 kernels are used in this block | |
| # By default, dropout is set to 0 | |
| self.block3 = custBlock(l_input_c=256, l_output_c=512, | |
| max_pooling=True, dropout= dropout, | |
| residual_links=2 | |
| ) # output_size = 4, rf_out = 57 | |
| self.max_pool_layer = nn.MaxPool2d(4, 4) | |
| # output_size = 1, rf_out = 81 | |
| self.flatten_layer = nn.Flatten() | |
| self.fc = nn.Linear(512, 10) | |
| #self.softmax = nn.Softmax() | |
| def forward(self, x): | |
| x = self.prep_block(x) | |
| x = self.block1(x) | |
| x = self.block2(x) | |
| x = self.block3(x) | |
| x = self.max_pool_layer(x) | |
| x = self.flatten_layer(x) | |
| x = self.fc(x) | |
| return x | |
| # Network Summary | |
| def summary(self, input_size=None, depth=10): | |
| return torchinfo.summary(self, input_size=input_size, | |
| depth=depth, | |
| col_names=["input_size", | |
| "output_size", | |
| "num_params", | |
| "kernel_size", | |
| "params_percent"]) | |