VisionEncoderDecoder is a local image-to-text transformer that joins a vision encoder to a text decoder. The encoder produces image embeddings; the decoder turns them into caption tokens. section 4.3.2
Trace one forward pass through three connected stages:
- Encode the image into a sequence of image embeddings. section 4.3.2
- Embed the input sequence, beginning with <start>, and apply masked self-attention so each prediction uses only preceding tokens. section 4.3.2
- Apply cross-attention between decoder states and image embeddings, then predict a probability distribution over the vocabulary. section 4.3.2 section 4.4.1
Training on a caption
For “a dog runs,” feed <start> a dog runs as the input and a dog runs <end> as the target. Compare predictions with the target using cross-entropy loss, then update the model parameters. section 4.3.2
Listing 4.7 The encoder–decoder transformer to add captions
class VisionEncoderDecoder(nn.Module):
def __init__(self, image_size, channels_in,
num_emb, patch_size=16,
hidden_size=128, num_layers=(6,6),
num_heads=4):
super(VisionEncoderDecoder, self).__init__()
self.encoder = VisionEncoder(
image_size=image_size, channels_in=channels_in,
patch_size=patch_size, hidden_size=hidden_size,
num_layers=num_layers[0], num_heads=num_heads)
self.decoder = Decoder(num_emb=num_emb,
hidden_size=hidden_size,
num_layers=num_layers[1], num_heads=num_heads)
def forward(self, input_image, target_seq, padding_mask):
bool_padding_mask = padding_mask == 0 #1
encoded_seq = self.encoder(image=input_image) #2
decoded_seq = self.decoder(input_seq=target_seq,
encoder_output=encoded_seq,
input_padding_mask=bool_padding_mask) #3
return decoded_seq
Inspect the class wiring that constructs the vision encoder and decoder and passes images, token sequences, and padding masks through the model. section 4.3.2
