JimmyChin1998/Pytorch-Learning-File

1"""
2Contains PyTorch model code to instantiate a TinyVGG model.
3"""
4import torch
5from torch import nn 
6
7class TinyVGG(nn.Module):
8  """Creates the TinyVGG architecture.
9
10  Replicates the TinyVGG architecture from the CNN explainer website in PyTorch.
11  See the original architecture here: https://poloclub.github.io/cnn-explainer/
12
13  Args:
14    input_shape: An integer indicating number of input channels.
15    hidden_units: An integer indicating number of hidden units between layers.
16    output_shape: An integer indicating number of output units.
17  """
18  def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
19      super().__init__()
20      self.conv_block_1 = nn.Sequential(
21          nn.Conv2d(in_channels=input_shape, 
22                    out_channels=hidden_units, 
23                    kernel_size=3, 
24                    stride=1, 
25                    padding=0),  
26          nn.ReLU(),
27          nn.Conv2d(in_channels=hidden_units, 
28                    out_channels=hidden_units,
29                    kernel_size=3,
30                    stride=1,
31                    padding=0),
32          nn.ReLU(),
33          nn.MaxPool2d(kernel_size=2,
34                        stride=2)
35      )
36      self.conv_block_2 = nn.Sequential(
37          nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),
38          nn.ReLU(),
39          nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),
40          nn.ReLU(),
41          nn.MaxPool2d(2)
42      )
43      self.classifier = nn.Sequential(
44          nn.Flatten(),
45          # Where did this in_features shape come from? 
46          # It's because each layer of our network compresses and changes the shape of our inputs data.
47          nn.Linear(in_features=hidden_units*13*13,
48                    out_features=output_shape)
49      )
50
51  def forward(self, x: torch.Tensor):
52      x = self.conv_block_1(x)
53      x = self.conv_block_2(x)
54      x = self.classifier(x)
55      return x
56      # return self.classifier(self.conv_block_2(self.conv_block_1(x))) # <- leverage the benefits of operator fusion
57