Pytorch

Examples of PyTorch code with best practices

• Simple embedding layer implementation in PyTorch with best practices

import math

from torch import Tensor, nn

class Embeddings(nn.Module):
    """
    Embedding layer for a given vocabulary size and model dimension.
    """

    def __init__(self, d_model: int, vocab: int) -> None:
		"""
		Initializes the Embeddings module.

		Parameters
		----------
		d_model : int
			The dimension of the model (embedding size)
		vocab : int
			The size of the vocabulary (number of unique tokens)
		"""
    
        super().__init__()
        
		# Layer that maps input token indices to dense vectors of size `d_model`
		self.lut: nn.Embedding = nn.Embedding(vocab, d_model)
        self.d_model: int = d_model

    def forward(self, x: Tensor) -> Tensor:
		"""
		Performs the forward pass of the Embeddings module.

		x : Tensor
			Input tensor containing token indices of shape (batch_size, sequence_length)

		Returns
		-------
		Tensor
			Output tensor containing the embedded representations of the input tokens,
			with shape (batch_size, sequence_length, d_model)
		"""

        out: Tensor = self.lut(x) * math.sqrt(self.d_model)
        
		return out