BillyCoder13/Multi_Class_Image_Classification

1import math2import torch.nn as nn3import torch.utils.model_zoo as model_zoo4import torch.optim as optim5from torchvision import transforms6import time7import matplotlib.pyplot as plt8 9 10model_urls = {11    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',12    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',13    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',14        'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',15        'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',16    }17 18 19 20class BasicBlock(nn.Module):21    """22    This is a basic block that contains two convolutional layers followed by23    a batch normalization layer and a ReLU activation function, where the skip24    connection is added before the second relu.25    ---26 27    - inplanes: { int } - The number of input channels.28    - planes: { int } - The number of output channels.29    - stride: { int } - The stride of convolutional layers.30    - downsample: { nn.Sequential } - A sequential of convolutional layers that fit the31        identity mapping to the desired output size.32    """33    expansion = 134 35    def __init__(self, inplanes, planes, stride=1, downsample=None):36        super(BasicBlock, self).__init__()37        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride,38                               padding=1, bias=False)39        self.bn1 = nn.BatchNorm2d(planes)40        self.relu = nn.ReLU(inplace=True)41 42        self.conv2 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride,43                               padding=1, bias=False)44        self.bn2 = nn.BatchNorm2d(planes)45        self.downsample = downsample46        self.stride = stride47 48    def forward(self, x):49        """50        This is the forward pass of the basic block where the input tensor x is passed51        through the first convolutional layer, batch normalization layer, and the ReLU52        activation function. The result is passed through the second convolutional layer,53        batch normalization layer, and the ReLU activation function. The result is then54        added to the identity mapping and passed through the ReLU activation function.55        """56        residual = x57 58        # Convolve with a 3X3Xplanes kernel59        out = self.conv1(x)60        out = self.bn1(out)61        out = self.relu(out)62 63        # Convolve with a 3X3Xplanes kernel64        out = self.conv2(out)65        out = self.bn2(out)66 67        # If the stride is not 1 or the number of input channels is not equal68        # to the number of output channels then we need to fit the identity69        # mapping to the desired output size by applying the downsample.70        if self.downsample is not None:71            residual = self.downsample(x)72 73        # Add the identity mapping to the output of the second convolutional layer.74        out += residual75        # Apply the ReLU activation function after the addition.76        out = self.relu(out)77 78        return out79    80 81 82class Bottleneck(nn.Module):83    """84    This class defines a bottle neck that fits the identity mapping to the desired85    output size before adding it to the output of the following layers.86    ---87    - inplanes: { int } - The number of input channels.88    - planes: { int } - The number of output channels.89    - stride: { int } - The stride of the second convolutional layer.90    - downsample: { nn.Sequential } - A sequential of convolutional layers that fit the91        identity mapping to the desired output size.92 93    The following layers are defined:94        - A 1x1 convolutional layer (self.conv1) with inplanes input channels and planes95        output channels is defined.96        - A batch normalization layer (self.bn1) is defined for the output of self.conv1.97        - A 3x3 convolutional layer (self.conv2) with planes input channels, planes output98        channels, and stride 'stride' is defined.99        - A batch normalization layer (self.bn2) is defined for the output of self.conv2.100        - A 1x1 convolutional layer (self.conv3) with planes input channels101        and planes * self.expansion output channels is defined.102        - A batch normalization layer (self.bn3) is defined for the output of self.conv3.103        - A ReLU activation function (self.relu) is defined.104    """105    expansion = 4106 107    def __init__(self, inplanes, planes, stride=1, downsample=None):108        super(Bottleneck, self).__init__()109        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)110        self.bn1 = nn.BatchNorm2d(planes)111 112        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,113                               stride=stride, padding=1, bias=False)114        self.bn2 = nn.BatchNorm2d(planes)115 116        self.conv3 = nn.Conv2d(117            planes, planes * self.expansion, kernel_size=1, bias=False)118        self.bn3 = nn.BatchNorm2d(planes * self.expansion)119        self.relu = nn.ReLU(inplace=True)120 121        self.downsample = downsample122        self.stride = stride123 124    def forward(self, x):125        """126            The Forward Pass127            ----------------128            Steps:129 130            - The input tensor x is saved as residual.131            - x is passed through self.conv1, self.bn1, and self.relu.132            - The result is passed through self.conv2, self.bn2, and self.relu.133            - The result is passed through self.conv3 and self.bn3.134 135            - If self.downsample is not None, residual is passed through self.downsample.136            - The output of the previous step is added to out.137            - The result is passed through self.relu.138            - The result is returned.139        """140        residual = x141        # Convolve with a 1X1Xplanes kernel142        out = self.conv1(x)143        out = self.bn1(out)144        out = self.relu(out)145 146        # Convolve with a 3X3Xplanes kernel147        out = self.conv2(out)148        out = self.bn2(out)149        out = self.relu(out)150 151        # Convolve with a 1X1Xplanes*expansion kernel152        out = self.conv3(out)153        out = self.bn3(out)154 155        # If the stride is not 1 or the number of input channels is not equal156        # to the number of output channels then we need to fit the identity157        # mapping to the desired output size by applying the downsample.158        if self.downsample is not None:159            residual = self.downsample(x)160 161        out += residual162        # Apply the ReLU activation function after the addition.163        out = self.relu(out)164 165        return out166    167 168class ResNet(nn.Module):169    """170    This is the ResNet class that is used in ResNet50, ResNet101, and ResNet152.171    """172    def __init__(self, block, layers, stride=None):173        self.inplanes = 64174        super(ResNet, self).__init__()175        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)176        self.bn1 = nn.BatchNorm2d(64)177        self.relu = nn.ReLU(inplace=True)178        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)179        self.layer1 = self._make_layer(block, 64, layers[0], stride=stride[0])180        self.layer2 = self._make_layer(block, 128, layers[1], stride=stride[1])181        self.layer3 = self._make_layer(block, 256, layers[2], stride=stride[2])182        self.layer4 = self._make_layer(block, 512, layers[3], stride=stride[3])183        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))184 185        self.fc = nn.Linear(512 * block.expansion, 1000)186 187        for m in self.modules():188            if isinstance(m, nn.Conv2d):189                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels190                m.weight.data.normal_(0, math.sqrt(2. / n))191            elif isinstance(m, nn.BatchNorm2d):192                m.weight.data.fill_(1)193                m.bias.data.zero_()194 195    def _make_layer(self, block, planes, blocks, stride=1):196        downsample = None197        if stride != 1 or self.inplanes != planes * block.expansion:198            downsample = nn.Sequential(199                nn.Conv2d(self.inplanes, planes * block.expansion,200                          kernel_size=1, stride=stride, bias=False),201                nn.BatchNorm2d(planes * block.expansion),202            )203 204        layers = []205        layers.append(block(self.inplanes, planes, stride, downsample))206        self.inplanes = planes * block.expansion207        for i in range(1, blocks):208            layers.append(block(self.inplanes, planes))209 210        return nn.Sequential(*layers)211 212    def forward(self, x):213        x = self.conv1(x)214        x = self.bn1(x)215        x = self.relu(x)216        x = self.maxpool(x)217 218        x = self.layer1(x)219        x = self.layer2(x)220        x = self.layer3(x)221        x = self.layer4(x)222 223        x = self.avgpool(x)224        x = x.view(x.size(0), -1)225        x = self.fc(x)226 227        return x228 229 230 231 232def resnet50(pretrained=False, stride=None, num_classes=200, **kwargs):233    """Constructs a ResNet-50 model.234 235    Args:236        pretrained (bool): If True, returns a model pre-trained on ImageNet237        :param pretrained:238        :param stride:239    """240    if stride is None:241        stride = [1, 2, 2, 1]242    model = ResNet(Bottleneck, [3, 4, 6, 3], stride=stride, **kwargs)243    if pretrained:244        model.load_state_dict(model_zoo.load_url(245            model_urls['resnet50']), strict=True)246    if num_classes != 200:247        model.fc = nn.Linear(512 * Bottleneck.expansion, num_classes)248    return model