Creating the full PatchEmbedding class

[Marcusrem]

Hello, a bit of code in here, but the code seems to work 💯 :

class PatchEmbedding(nn.Module):
  """
  Transforms 2D image into a sequence of learnable embedding patches
  """
  def __init__(self,
               in_channels:int= 3,
               patch_size:int = 16,
               embedding_dim:int = 768,
               batch_size:int=1,
               number_of_patches:int = 196): # From table 1 for ViT-Base
    super().__init__()

    self.patch_size = patch_size
     # 3. Create the Conv2d layer with hyperparameters from the ViT paper to tunr the images into embedded patches
    self.patcher = nn.Conv2d(in_channels=in_channels, # number of color channels
                   out_channels= embedding_dim, # from Table 1: Hidden size D, this is the embedding size / embedding_layer_output_shape
                   kernel_size=patch_size, # could also use (patch_size, patch_size)
                   stride=patch_size,
                   padding=0)

    # 4. Create flatten layer to flatten feature map outputs of Conv2d
    self.flatten = nn.Flatten(start_dim=2, # flatten feature_map_height (dimension 2)
                         end_dim=3) # flatten feature_map_width (dimension 3)

    # 5. Create a class token embedding
    self.class_token = nn.Parameter(torch.ones(batch_size, 1,embedding_dim), 
                                    requires_grad=True)

    # 6. Create a positional embedding
    self.positional_embedding = nn.Parameter(torch.ones(batch_size,number_of_patches+1,embedding_dim),
                                             requires_grad=True)


  # 7. Define a forward method to define the forward computation steps
  def forward(self,x):
    # Create an assertion to check that inputs are the correct shape
    image_resolution = x.shape[-1]
    assert image_resolution % patch_size == 0, f"Input size must be divisible by patch size, image shape:{image_resolution} patch size:{patch_size}"

    # Perform the forward pass
    x_patched = self.patcher(x)
    x_flattened = self.flatten(x_patched)
    x_patch_embeding_class_token = torch.cat((self.class_token,x_flattened.permute(0,2,1)),dim=1)
    x_patch_position_embedding = self.positional_embedding+x_patch_embeding_class_token
    # 6. Make the returned sequence embedding dim are in the right order(batch_size, number_of_patches, embedding_dim)
    return x_patch_position_embedding
    ```

[Marcusrem]

A great course anyway @mrdbourke!
What do you think, it's ok in this sec?
Greetings!!