| |
| |
| |
| import torch |
| from torch import nn |
|
|
|
|
| |
| |
| |
| class VGGBlock(nn.Module): |
| ''' |
| Defines a modified block in the VGG architecture, |
| which includes batch normalization between convolutional layers and ReLU activations. |
| |
| Reference: https://poloclub.github.io/cnn-explainer/ |
| Reference: https://d2l.ai/chapter_convolutional-modern/vgg.html |
| |
| Args: |
| num_convs (int): Number of consecutive convolutional layers + ReLU activations. |
| in_channels (int): Number of channels in the input. |
| hidden_channels (int): Number of hidden channels between convolutional layers. |
| out_channels (int): Number of channels in the output. |
| ''' |
| def __init__(self, |
| num_convs: int, |
| in_channels: int, |
| hidden_channels: int, |
| out_channels: int): |
| super().__init__() |
| |
| self.layers = [] |
| |
| for i in range(num_convs): |
| conv_in = in_channels if i == 0 else hidden_channels |
| conv_out = out_channels if i == num_convs-1 else hidden_channels |
| |
| self.layers += [ |
| nn.Conv2d(conv_in, conv_out, kernel_size = 3, stride = 1, padding = 1), |
| nn.BatchNorm2d(conv_out), |
| nn.ReLU() |
| ] |
| |
| self.layers.append(nn.MaxPool2d(kernel_size = 2, stride = 2)) |
| |
| self.vgg_blk = nn.Sequential(*self.layers) |
| |
| def forward(self, X: torch.Tensor) -> torch.Tensor: |
| ''' |
| Forward pass of VGG block. |
| |
| Args: |
| X (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width) |
| Returns: |
| torch.Tensor: Output tensor of shape (batch_size, out_channels, new_height, new_width) |
| ''' |
| |
| return self.vgg_blk(X) |
| |
| class TinyVGG(nn.Module): |
| ''' |
| Creates a simplified version of a VGG model, adapted from |
| https://github.com/poloclub/cnn-explainer/blob/master/tiny-vgg/tiny-vgg.py. |
| The main difference is that the hidden dimensions and number of convolutional layers |
| remain the same across VGG blocks and the classifier's linear layers has output fewer features. |
| |
| Args: |
| num_blks (int): Number of VGG blocks to put in the model |
| num_convs (int): Number of consecutive convolutional layers + ReLU activations in each VGG block. |
| in_channels (int): Number of channels in the input. |
| hidden_channels (int): Number of hidden channels between convolutional layers. |
| fc_hidden_dim (int): Number of output (hidden) features for the first linear layer of the classifer. |
| num_classes (int): Number of class labels. |
| |
| ''' |
| def __init__(self, |
| num_blks: int, |
| num_convs: int, |
| in_channels: int, |
| hidden_channels: int, |
| fc_hidden_dim: int, |
| num_classes: int): |
| super().__init__() |
| |
| self.all_blks = [] |
| for i in range(num_blks): |
| conv_in = in_channels if i == 0 else hidden_channels |
| self.all_blks.append( |
| VGGBlock(num_convs, conv_in, hidden_channels, hidden_channels) |
| ) |
| |
| self.vgg_body = nn.Sequential(*self.all_blks) |
| self.classifier = nn.Sequential( |
| nn.Flatten(), |
| nn.LazyLinear(fc_hidden_dim), nn.ReLU(), nn.Dropout(0.5), |
| nn.LazyLinear(num_classes) |
| ) |
| |
| self.vgg_body.apply(self._custom_init) |
| self.classifier.apply(self._custom_init) |
| |
| def _custom_init(self, module): |
| ''' |
| Initializes convolutional layer weights with Xavier initialization method. |
| Initializes convolutional layer biases to zero. |
| ''' |
| if isinstance(module, (nn.Conv2d)): |
| nn.init.xavier_uniform_(module.weight) |
| nn.init.zeros_(module.bias) |
| |
| def forward(self, X: torch.Tensor) -> torch.Tensor: |
| ''' |
| Forward pass of the TinyVGG model. |
| |
| Args: |
| X (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width). |
| |
| Returns: |
| torch.Tensor: Logits of shape (batch_size, num_classes). |
| ''' |
| |
| X = self.vgg_body(X) |
| return self.classifier(X) |
|
|
