Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import math | |
| class ExtractPatches(nn.Module): | |
| def __init__(self, patch_size: int = 16): | |
| super().__init__() | |
| self.patch_size = patch_size | |
| self.unfold = nn.Unfold(kernel_size=patch_size, stride=patch_size) | |
| def forward(self, x): | |
| batch_size, c, h, w = x.shape | |
| # Unfold applies a slding window to generate patches | |
| # The transpose and reshape change the shape to (batch_size, num_patches, 3 * patch_size * patch_size), flattening the patches | |
| return ( | |
| self.unfold(x) | |
| .transpose(1, 2) | |
| .reshape(batch_size, -1, c * self.patch_size * self.patch_size) | |
| ) | |
| # Positional Encoding | |
| class PositionalEncoding(nn.Module): | |
| def __init__(self, d_model: int): | |
| """ | |
| d_model: dimensions of the embeddings (number of values in each embedding vector) | |
| """ | |
| super().__init__() | |
| # Intead of precomputing fixed values, we will compute in the forward pass based off of the sinusodiual encoding formula | |
| self.d_model = d_model | |
| def forward(self, x): | |
| device = x.device | |
| half_dim = self.d_model // 2 # Use half for sin and half for cos | |
| emb = math.log(10000.0) / (half_dim - 1) | |
| emb = torch.exp(torch.arange(half_dim, device=device) * -emb) | |
| emb = x[:, None] * emb[None, :] # (batch_size, half_dim) | |
| emb = torch.cat((emb.sin(), emb.cos()), dim=-1) | |
| return emb | |
| # Multi-Head Self-Attention | |
| class MultiHeadAttention(nn.Module): | |
| def __init__(self, d_model: int = 512, n_heads: int = 8, dropout: float = 0.1): | |
| """ | |
| d_model: dimensions of the embeddings (number of values in each embedding vector) | |
| n_heads: number of self attention heads per sequence | |
| dropout: probability of dropout | |
| """ | |
| super().__init__() | |
| assert ( | |
| d_model % n_heads == 0 | |
| ) # We want to make sure that the dimensions are split evenly among the attention heads. | |
| self.d_model = d_model | |
| self.n_heads = n_heads | |
| self.d_key = d_model // n_heads | |
| self.Wq = nn.Linear(d_model, d_model) # Learnable weights for query | |
| self.Wk = nn.Linear(d_model, d_model) # Learnable weights for key | |
| self.Wv = nn.Linear(d_model, d_model) # Learnable weights for value | |
| self.Wo = nn.Linear(d_model, d_model) # Learnable weights for output | |
| self.dropout = nn.Dropout(p=dropout) | |
| def forward(self, query, key, value, mask=None): | |
| """ | |
| query: (batch_size, q_length, d_model) | |
| key: (batch_size, k_length, d_model) | |
| value: (batch_size, s_length, d_model) | |
| """ | |
| batch_size = key.size(0) | |
| # Matrix multiplication for Q, K, and V tensors | |
| Q = self.Wq(query) | |
| K = self.Wk(key) | |
| V = self.Wv(value) | |
| # Split each tensor into heads | |
| Q = Q.view(batch_size, -1, self.n_heads, self.d_key).permute( | |
| 0, 2, 1, 3 | |
| ) # (batch_size, n_heads, q_length, d_key) | |
| K = K.view(batch_size, -1, self.n_heads, self.d_key).permute( | |
| 0, 2, 1, 3 | |
| ) # (batch_size, n_heads, k_length, d_key) | |
| V = V.view(batch_size, -1, self.n_heads, self.d_key).permute( | |
| 0, 2, 1, 3 | |
| ) # (batch_size, n_heads, v_length, d_key) | |
| # Scaled dot product | |
| # K^T becomees (batch_size, n_heads, d_key, k_length) | |
| scaled_dot_product = torch.matmul(Q, K.permute(0, 1, 3, 2)) / math.sqrt( | |
| self.d_key | |
| ) # (batch_size, n_heads, q_length, k_length) | |
| if mask is not None: | |
| scaled_dot_product = scaled_dot_product.masked_fill( | |
| mask == 0, -float("inf") | |
| ) # Filling it with 0 would result in 1 after the mask because e^0 = 1. Intead we fill it with an infinitley large negative number | |
| # Softmax function for attention probabilities | |
| attention_probs = torch.softmax(scaled_dot_product, dim=-1) | |
| # Multiply by V to get attention with respect to the values | |
| A = torch.matmul(self.dropout(attention_probs), V) | |
| # Reshape attention back to (batch_size, q_length, d_model) | |
| A = ( | |
| A.permute(0, 2, 1, 3) | |
| .contiguous() | |
| .view(batch_size, -1, self.n_heads * self.d_key) | |
| ) | |
| # Pass through the final linear layer | |
| output = self.Wo(A) | |
| return ( | |
| output, | |
| attention_probs, | |
| ) # Output shape: (batch_size, q_length, d_model), Attention probs shape: (batch_size, n_heads, q_length, k_length) | |
| # Position-Wise Feed Forward Network (FFN) | |
| class PositionwiseFeedForward(nn.Module): | |
| def __init__(self, d_model: int, dropout: float = 0.1): | |
| """ | |
| d_model: dimensions of the embeddings (number of values in each embedding vector) | |
| dropout: probability of dropout | |
| """ | |
| super().__init__() | |
| self.ffn = nn.Sequential( | |
| nn.Linear(in_features=d_model, out_features=(d_model * 4)), | |
| nn.ReLU(), | |
| nn.Linear(in_features=(d_model * 4), out_features=d_model), | |
| nn.Dropout(p=dropout), | |
| ) | |
| def forward(self, x): | |
| return self.ffn(x) | |
| # Encoder Layer | |
| class EncoderLayer(nn.Module): | |
| def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1): | |
| """ | |
| d_model: dimensions of the embeddings (number of values in each embedding vector) | |
| n_heads: number of self attention heads per sequence | |
| dropout: probability of dropout | |
| """ | |
| super().__init__() | |
| # Multi-Head Self-Attention sublayer | |
| self.attention = MultiHeadAttention( | |
| d_model=d_model, n_heads=n_heads, dropout=dropout | |
| ) | |
| self.attention_layer_norm = nn.LayerNorm(d_model) # Layer normalization | |
| # Position-wise Feed-forward Network | |
| self.position_wise_ffn = PositionwiseFeedForward( | |
| d_model=d_model, dropout=dropout | |
| ) | |
| self.ffn_layer_norm = nn.LayerNorm(d_model) # Layer normalization | |
| self.dropout = nn.Dropout(p=dropout) | |
| def forward(self, src): | |
| """ | |
| src: embedded sequences (batch_size, seq_length, d_model) | |
| """ | |
| # Multi-Head Attention | |
| _src, attention_probs = self.attention( | |
| src, src, src, None | |
| ) # Q, K, V, src_mask: we don't need a source mask because all images are the same dimension | |
| # Residual Addition and Layer Normalization | |
| src = self.attention_layer_norm( | |
| src + self.dropout(_src) | |
| ) # We do residual addition by adding back the src (the embeddings) to the output of Self-Attention | |
| # Position-wise Feed-forward Network | |
| _src = self.position_wise_ffn(src) | |
| # Residual Addition and Layer Normalization | |
| src = self.ffn_layer_norm(src + self.dropout(_src)) | |
| return src, attention_probs | |
| # The Encoder that takes in images and returns the encoding to be passed into the decoder | |
| class Encoder(nn.Module): | |
| def __init__( | |
| self, | |
| image_size: int, | |
| in_channels: int, | |
| patch_size: int = 16, | |
| d_model: int = 128, | |
| n_layers: int = 3, | |
| n_heads: int = 4, | |
| dropout: float = 0.1, | |
| ): | |
| """ | |
| d_model: dimensions of the embeddings (number of values in each embedding vector) | |
| n_layers: number of encoder layers in the encoder block | |
| n_heads: number of self attention heads per sequence | |
| dropout: probability of dropout | |
| """ | |
| super().__init__() | |
| self.patch_size = patch_size | |
| self.extract_patches = ExtractPatches(patch_size=patch_size) | |
| self.fc_in = nn.Linear(in_channels * patch_size * patch_size, d_model) | |
| seq_length = (image_size // patch_size) ** 2 | |
| # Image src is going to use a learnable positional encoding | |
| self.pos_embedding = nn.Parameter( | |
| torch.empty(1, seq_length, d_model).normal_(std=0.02) | |
| ) | |
| # Create n_layers encoders | |
| self.layers = nn.ModuleList( | |
| [ | |
| EncoderLayer(d_model=d_model, n_heads=n_heads, dropout=dropout) | |
| for layer in range(n_layers) | |
| ] | |
| ) | |
| self.dropout = nn.Dropout(p=dropout) | |
| def forward(self, src): | |
| """ | |
| src: embedded sequences (batch_size, seq_length, d_model) | |
| """ | |
| # Extract the patches and apply a linear layer | |
| batch_size = src.shape[0] | |
| src = self.fc_in(self.extract_patches(src)) | |
| # Add the learned positional embedding | |
| src = src + self.pos_embedding | |
| # Pass the sequences through each encoder layer | |
| for layer in self.layers: | |
| src, attention_probs = layer(src) | |
| self.attention_probs = attention_probs | |
| return src | |
| # Decoder Layer | |
| class DecoderLayer(nn.Module): | |
| def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1): | |
| """ | |
| d_model: dimensions of the embeddings (number of values in each embedding vector) | |
| n_heads: number of self attention heads per sequence | |
| dropout: probability of dropout | |
| """ | |
| super().__init__() | |
| # Masked Multi-Head Self-Attention sublayer | |
| self.masked_attention = MultiHeadAttention( | |
| d_model=d_model, n_heads=n_heads, dropout=dropout | |
| ) | |
| self.masked_attention_layer_norm = nn.LayerNorm(d_model) # Layer normalization | |
| # Multi-Head Self-Attention sublayer | |
| self.attention = MultiHeadAttention( | |
| d_model=d_model, n_heads=n_heads, dropout=dropout | |
| ) | |
| self.attention_layer_norm = nn.LayerNorm(d_model) # Layer normalization | |
| # Position-wise Feed-forward Network | |
| self.position_wise_ffn = PositionwiseFeedForward( | |
| d_model=d_model, dropout=dropout | |
| ) | |
| self.ffn_layer_norm = nn.LayerNorm(d_model) # Layer normalization | |
| self.dropout = nn.Dropout(p=dropout) | |
| def forward(self, trg, src, trg_mask): | |
| """ | |
| trg: embedded captions (batch_size, trg_seq_length, d_model) | |
| src: embedded images (batch_size, src_seq_length, d_model) | |
| trg_mask: mask for the captions preventing peeking at future tokens (batch_size, 1, trg_seq_length, trg_seq_length) | |
| """ | |
| # Masked Multi-Head Attention | |
| # The target mask is used to prevent the model from seeing future tokens. This ensures that the prediction is made solely based on past and present tokens. | |
| _trg, masked_attention_probs = self.masked_attention( | |
| trg, trg, trg, trg_mask | |
| ) # Q, K, V, mask | |
| # Residual Addition and Layer Normalization | |
| trg = self.masked_attention_layer_norm(trg + self.dropout(_trg)) | |
| # Multi-Head Attention - This time, we also pass in the output of the encoder layers as src. | |
| # This is important because this allows us to keep track of and learn relationships between the input and output tokens. | |
| _trg, attention_probs = self.attention(trg, src, src, None) # Q, K, V, mask | |
| # Residual Addition and Layer Normalization | |
| trg = self.attention_layer_norm(trg + self.dropout(_trg)) | |
| # Position-wise Feed-forward Network | |
| _trg = self.position_wise_ffn(trg) | |
| # Residual Addition and Layer Normalization | |
| trg = self.ffn_layer_norm(trg + self.dropout(_trg)) | |
| return trg, attention_probs, masked_attention_probs | |
| # The Decoder Module that takes the encoded images from the encoder and generates captions | |
| class Decoder(nn.Module): | |
| def __init__( | |
| self, | |
| vocab_size: int, | |
| d_model: int = 128, | |
| n_layers: int = 3, | |
| n_heads: int = 4, | |
| dropout: float = 0.1, | |
| ): | |
| """ | |
| vocab_size: size of the target vocabulary | |
| d_model: dimensions of the embeddings (number of values in each embedding vector) | |
| n_layers: number of encoder layers in the encoder block | |
| n_heads: number of self attention heads per sequence | |
| dropout: probability of dropout | |
| """ | |
| super().__init__() | |
| self.embedding = nn.Embedding(vocab_size, d_model) | |
| self.embedding.weight.data = 0.001 * self.embedding.weight.data | |
| # Initialize sinusoidal positional embeddings | |
| self.pos_emb = PositionalEncoding(d_model=d_model) | |
| # Create n_layers decoders | |
| self.layers = nn.ModuleList( | |
| [ | |
| DecoderLayer(d_model=d_model, n_heads=n_heads, dropout=dropout) | |
| for layer in range(n_layers) | |
| ] | |
| ) | |
| self.dropout = nn.Dropout(p=dropout) | |
| # Output layer | |
| self.Wo = nn.Linear(in_features=d_model, out_features=vocab_size) | |
| def make_trg_mask(self, trg): | |
| seq_length = trg.shape[1] | |
| trg_mask = torch.tril( | |
| torch.ones((seq_length, seq_length), device=trg.device) | |
| ).bool() | |
| return trg_mask.unsqueeze(0).unsqueeze( | |
| 0 | |
| ) # (batch_size=1, n_heads=1, seq_length, seq_length) | |
| def forward(self, trg, src): | |
| """ | |
| trg: target sequences (batch_size, trg_seq_length, d_model) | |
| src: embedding images (batch_size, src_seq_length, d_model) | |
| """ | |
| # Embed the target captions | |
| trg = self.embedding(trg) | |
| batch_size, l, h = trg.shape | |
| trg_index = torch.arange(l, device=trg.device) | |
| pos_emb = self.pos_emb(trg_index).reshape(1, l, h).expand(batch_size, l, h) | |
| # Add the fixed sinusodial positional embedding | |
| trg += pos_emb | |
| # Create a target mask for the target captions to prevent the model from peeking at future tokens | |
| trg_mask = self.make_trg_mask( | |
| trg | |
| ) # (batch_size, 1, trg_seq_length, trg_seq_length) | |
| # Pass the sequences through each decoder layer | |
| for layer in self.layers: | |
| trg, attention_probs, masked_attention_probs = layer(trg, src, trg_mask) | |
| self.attention_probs = attention_probs | |
| self.masked_attention_probs = masked_attention_probs # (batch_size, n_heads, trg_seq_len, src_seq_len) trg_seq_len: length of the target caption \ src_seq_len: number of patches from the encoder | |
| # Final linear output layer | |
| return self.Wo(trg) | |
| class CaptioningTransformer(nn.Module): | |
| def __init__( | |
| self, | |
| image_size: int, | |
| in_channels: int, | |
| vocab_size: int, | |
| device, | |
| patch_size: int = 16, | |
| d_model: int = 128, | |
| n_layers: int = 3, | |
| n_heads: int = 4, | |
| ): | |
| super().__init__() | |
| self.device = device | |
| # Create an encoder and decoder with specified parameters | |
| self.encoder = Encoder( | |
| image_size=image_size, | |
| in_channels=in_channels, | |
| patch_size=patch_size, | |
| d_model=d_model, | |
| n_layers=n_layers, | |
| n_heads=n_heads, | |
| ) | |
| self.decoder = Decoder( | |
| vocab_size=vocab_size, d_model=d_model, n_layers=n_layers, n_heads=n_heads | |
| ) | |
| def forward(self, src, trg): | |
| # Encoder layers | |
| src = self.encoder(src) # (batch_size, src_seq_length, d_model) | |
| # Decoder layers | |
| output = self.decoder( | |
| trg, src | |
| ) # Pass in both the target (for Masked Multi-Head Self-Attention) and source for (Cross-Attention) | |
| return output | |