[Chapter 8] Paper Replicating - Embedding class for image, class and position

[jsnellBioSig]

Helo Daniel,

Thank you for your ML classes, I have completed your TensorFlow corse and am working through the PyTorch class - both are easy to follow (as easy as learning ML can be :) and I am happy to say that virtually all code works as advertised. I am running on Mac (intel w/AMD GPU and M3) and am able to work through all notebooks with both using Apple Metal.

I took a stab at your challenge to create a class for image, class and position embeddings (at end of section 4) - this is what I came up with:

class PatchEmbedding2(nn.Module):
    """Turns a 2D input image into a 1D sequence learnable embedding vector.

    Args:
        in_channels (int): Number of color channels for the input images. Defaults to 3.
        patch_size (int): Size of patches to convert input image into. Defaults to 16.
        embedding_dim (int): Size of embedding to turn image into. Defaults to 768.
    """
    # 2. Initialize the class with appropriate variables
    def __init__(self,
                 in_channels:int=3,
                 patch_size:int=16,
                 embedding_dim:int=768,
                 batch_size:int=1,
                 number_of_patches:int=196):
        super().__init__()

        # 3. Create a layer to turn an image into patches
        self.patcher = nn.Conv2d(in_channels=in_channels,
                                 out_channels=embedding_dim,
                                 kernel_size=patch_size,
                                 stride=patch_size,
                                 padding=0)

        # 4. Create a layer to flatten the patch feature maps into a single dimension
        self.flatten = nn.Flatten(start_dim=2, # only flatten the feature map dimensions into a single vector
                                  end_dim=3)

        # 5. Create class token embedding
        self.class_token = nn.Parameter(torch.ones(batch_size, embedding_dim, 1),
                                   requires_grad=True) # make sure it's learnable
        
        # 8. Create position embedding
        self.position_embedding = nn.Parameter(torch.ones(1, embedding_dimension, number_of_patches+1),
                                          requires_grad=True) # make sure it's learnable

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

        # Perform the forward pass
        x_patched = self.patcher(x) # Create embedding patches
        x_flattened = self.flatten(x_patched) # Flatten patch embedding
        # print(f"class_token shape: {self.class_token.shape}")
        # print(f"x_flattened shape: {x_flattened.shape}")
        x_with_class = torch.cat((self.class_token, x_flattened),dim=2) # Prepend class token embedding to patch embedding
        # print(f"self.position_embedding shape: {self.position_embedding.shape}")
        # print(f"x_with_class shape: {x_with_class.shape}")
        x_with_class_and_pos = x_with_class + self.position_embedding # Add position embedding to patch embedding with class token
      
        # 10. Make sure the output shape has the right order
        return x_with_class_and_pos.permute(0, 2, 1) # adjust so the embedding is on the final dimension [batch_size, P^2•C, N] -> [batch_size, N, P^2•C]

set_seeds()

# Create an instance of patch embedding layer
patchify2 = PatchEmbedding2(in_channels=3,
                          patch_size=16,
                          embedding_dim=768,
                          batch_size=1,
                          number_of_patches=196)

# Pass a single image through
print(f"Input image shape: {image.unsqueeze(0).shape}")
patch_embedded_image2 = patchify2(image.unsqueeze(0)) # add an extra batch dimension on the 0th index, otherwise will error
print(f"Output patch embedding shape: {patch_embedded_image2.shape}")

print(f"\nNEWER patch and position embedding shape: {patch_embedded_image2.shape}")
print(patch_embedded_image2)

It seems to work - I welcome any comment you may have.

ALSO,

I am interested in playing around with a customized Object Detection - with object location (there seem to be a few different formats around). I would have many objects per image - would like to train on custom data,I would have a small-ish number of classes but not typical scenes from typical pretrained models... just curious what your suggestion might be as to where to start?? There are many potential choices it seems...

Thank you - and best regards,

Jeff