import math import torch.nn as nn import torch.utils.model_zoo as model_zoo import torch.optim as optim from torchvision import transforms import time import matplotlib.pyplot as plt model_urls = { 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', } class BasicBlock(nn.Module): """ This is a basic block that contains two convolutional layers followed by a batch normalization layer and a ReLU activation function, where the skip connection is added before the second relu. --- - inplanes: { int } - The number of input channels. - planes: { int } - The number of output channels. - stride: { int } - The stride of convolutional layers. - downsample: { nn.Sequential } - A sequential of convolutional layers that fit the identity mapping to the desired output size. """ expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None): super(BasicBlock, self).__init__() self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(planes) self.relu = nn.ReLU(inplace=True) self.conv2 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(planes) self.downsample = downsample self.stride = stride def forward(self, x): """ This is the forward pass of the basic block where the input tensor x is passed through the first convolutional layer, batch normalization layer, and the ReLU activation function. The result is passed through the second convolutional layer, batch normalization layer, and the ReLU activation function. The result is then added to the identity mapping and passed through the ReLU activation function. """ residual = x # Convolve with a 3X3Xplanes kernel out = self.conv1(x) out = self.bn1(out) out = self.relu(out) # Convolve with a 3X3Xplanes kernel out = self.conv2(out) out = self.bn2(out) # If the stride is not 1 or the number of input channels is not equal # to the number of output channels then we need to fit the identity # mapping to the desired output size by applying the downsample. if self.downsample is not None: residual = self.downsample(x) # Add the identity mapping to the output of the second convolutional layer. out += residual # Apply the ReLU activation function after the addition. out = self.relu(out) return out class Bottleneck(nn.Module): """ This class defines a bottle neck that fits the identity mapping to the desired output size before adding it to the output of the following layers. --- - inplanes: { int } - The number of input channels. - planes: { int } - The number of output channels. - stride: { int } - The stride of the second convolutional layer. - downsample: { nn.Sequential } - A sequential of convolutional layers that fit the identity mapping to the desired output size. The following layers are defined: - A 1x1 convolutional layer (self.conv1) with inplanes input channels and planes output channels is defined. - A batch normalization layer (self.bn1) is defined for the output of self.conv1. - A 3x3 convolutional layer (self.conv2) with planes input channels, planes output channels, and stride 'stride' is defined. - A batch normalization layer (self.bn2) is defined for the output of self.conv2. - A 1x1 convolutional layer (self.conv3) with planes input channels and planes * self.expansion output channels is defined. - A batch normalization layer (self.bn3) is defined for the output of self.conv3. - A ReLU activation function (self.relu) is defined. """ expansion = 4 def __init__(self, inplanes, planes, stride=1, downsample=None): super(Bottleneck, self).__init__() self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) self.bn1 = nn.BatchNorm2d(planes) self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(planes) self.conv3 = nn.Conv2d( planes, planes * self.expansion, kernel_size=1, bias=False) self.bn3 = nn.BatchNorm2d(planes * self.expansion) self.relu = nn.ReLU(inplace=True) self.downsample = downsample self.stride = stride def forward(self, x): """ The Forward Pass ---------------- Steps: - The input tensor x is saved as residual. - x is passed through self.conv1, self.bn1, and self.relu. - The result is passed through self.conv2, self.bn2, and self.relu. - The result is passed through self.conv3 and self.bn3. - If self.downsample is not None, residual is passed through self.downsample. - The output of the previous step is added to out. - The result is passed through self.relu. - The result is returned. """ residual = x # Convolve with a 1X1Xplanes kernel out = self.conv1(x) out = self.bn1(out) out = self.relu(out) # Convolve with a 3X3Xplanes kernel out = self.conv2(out) out = self.bn2(out) out = self.relu(out) # Convolve with a 1X1Xplanes*expansion kernel out = self.conv3(out) out = self.bn3(out) # If the stride is not 1 or the number of input channels is not equal # to the number of output channels then we need to fit the identity # mapping to the desired output size by applying the downsample. if self.downsample is not None: residual = self.downsample(x) out += residual # Apply the ReLU activation function after the addition. out = self.relu(out) return out class ResNet(nn.Module): """ This is the ResNet class that is used in ResNet50, ResNet101, and ResNet152. """ def __init__(self, block, layers, stride=None): self.inplanes = 64 super(ResNet, self).__init__() self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = nn.BatchNorm2d(64) self.relu = nn.ReLU(inplace=True) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) self.layer1 = self._make_layer(block, 64, layers[0], stride=stride[0]) self.layer2 = self._make_layer(block, 128, layers[1], stride=stride[1]) self.layer3 = self._make_layer(block, 256, layers[2], stride=stride[2]) self.layer4 = self._make_layer(block, 512, layers[3], stride=stride[3]) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) self.fc = nn.Linear(512 * block.expansion, 1000) for m in self.modules(): if isinstance(m, nn.Conv2d): n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels m.weight.data.normal_(0, math.sqrt(2. / n)) elif isinstance(m, nn.BatchNorm2d): m.weight.data.fill_(1) m.bias.data.zero_() def _make_layer(self, block, planes, blocks, stride=1): downsample = None if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(planes * block.expansion), ) layers = [] layers.append(block(self.inplanes, planes, stride, downsample)) self.inplanes = planes * block.expansion for i in range(1, blocks): layers.append(block(self.inplanes, planes)) return nn.Sequential(*layers) def forward(self, x): x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x = self.maxpool(x) x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) x = self.avgpool(x) x = x.view(x.size(0), -1) x = self.fc(x) return x def resnet50(pretrained=False, stride=None, num_classes=200, **kwargs): """Constructs a ResNet-50 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet :param pretrained: :param stride: """ if stride is None: stride = [1, 2, 2, 1] model = ResNet(Bottleneck, [3, 4, 6, 3], stride=stride, **kwargs) if pretrained: model.load_state_dict(model_zoo.load_url( model_urls['resnet50']), strict=True) if num_classes != 1000: model.fc = nn.Linear(512 * Bottleneck.expansion, num_classes) return model