Resource retrieval

Get the pre-trained net:

In[1]:=

Out[1]=

NetModel parameters

This model consists of a family of individual nets, each identified by a specific parameter combination. Inspect the available parameters:

In[2]:=

Out[2]=

Pick a non-default net by specifying the parameters:

In[3]:=

Out[3]=

Pick a non-default uninitialized net:

In[4]:=

Out[4]=

Evaluation function

Define an evaluation function that uses all model parts to obtain the image features and automate the caption generation:

In[5]:=

Basic usage

Define a test image:

In[6]:=

Generate an image caption:

In[7]:=

Out[7]=

Define a test video and visualize the first five frames:

In[8]:=

Out[9]=

For zero-shot video-text retrieval, the mean image embedding is computed across uniformly sampled video frames:

In[10]:=

Out[10]=

Efficient generation

The evaluation function defined in the previous section is inefficient because every time a token of the caption is generated, the entire caption is read by the text encoder from the start. It is possible to avoid this by using NetUnfold. Reimplement the evaluation function using nets prepared by NetUnfold:

In[11]:=

In[12]:=

Prepare the unfolded net beforehand for the efficient image caption generation:

In[13]:=

The new textDecoder is an unfolded version of the old text encoder and text decoder combined and exposes ports for the hidden attention states:

In[14]:=

Out[14]=

In[15]:=

Out[15]=

Compare evaluation timing of the methods:

In[16]:=

In[17]:=

Out[17]=

In[18]:=

Out[18]=

The speedup in the efficient case depends on the length of the generated sequence, and image captioning models generally don't generate very long pieces of text. Hence, it's not very significant.

Feature space visualization

Get a set of images:

In[19]:=

In[20]:=

Out[20]=

Define a list of sentences in two categories:

In[21]:=

In[22]:=

Out[22]=

Connecting text and images

Define a test image:

In[23]:=

Define a list of text descriptions:

In[24]:=

Embed the test image and text descriptions into the same feature space:

In[25]:=

Rank the text description with respect to the correspondence to the input image according to the CosineDistance. Smaller distances mean higher correspondence between the text and the image:

In[26]:=

Out[26]=

Zero-shot image classification

By using the text and image feature extractors together, it's possible to perform generic image classification between any set of classes without having to explicitly train any model for those particular classes (zero-shot classification). Obtain the FashionMNIST test data, which contains ten thousand test images and 10 classes:

In[27]:=

Display a few random examples from the set:

In[28]:=

Out[28]=

Get a mapping between class IDs and labels:

In[29]:=

Out[29]=

Generate the text templates for the FashionMNIST labels and embed them. The text templates will effectively act as classification labels:

In[30]:=

Out[30]=

In[31]:=

In[32]:=

Out[32]=

Classify an image from the test set. Obtain its embedding:

In[33]:=

Out[33]=

In[34]:=

In[35]:=

Out[35]=

The result of the classification is the description of the embedding that is closest to the image embedding:

In[36]:=

Out[36]=

Find the top 10 descriptions nearest to the image embedding:

In[37]:=

Out[37]=

Obtain the accuracy of this procedure on the entire test set. Extract the features for all the images (if a GPU is available, setting TargetDevice -> "GPU" is recommended as the computation will take several minutes on CPU):

In[38]:=

In[39]:=

Out[39]=

Calculate the distance matrix between the computed text and image embeddings:

In[40]:=

Obtain the top-1 prediction:

In[41]:=

Obtain the final classification results:

In[42]:=

Out[42]=

Attention visualization for images

Just like the original Vision Transformer (see the model "Vision Transformer Trained on ImageNet Competition Data"), the image feature extractor divides the input images in 7x7 patches and performs self-attention on a set of 50 vectors: 49 vectors, or "tokens," representing the 7x7 patches and an additional one, a "feature extraction token," that is eventually used to produce the final feature representation of the image. Thus the attention procedure for this model can be visualized by inspecting the attention weights between the feature extraction token and the patch tokens. Define a test image:

In[43]:=

Extract the attention weights used for the last block of self-attention:

In[44]:=

In[45]:=

Out[45]=

Extract the attention weights between the feature extraction token and the input patches. These weights can be interpreted as which patches in the original image the net is "looking at" in order to perform the feature extraction:

In[46]:=

Out[47]=

Reshape the weights as a 3D array of 12 7x7 matrices. Each matrix corresponds to an attention head, while each element of the matrices corresponds to a patch in the original image:

In[48]:=

In[49]:=

Out[49]=

Visualize the attention weight matrices. Patches with higher values (red) are what is mostly being "looked at" for each attention head:

In[50]:=

Out[50]=

Define a function to visualize the attention matrix on an image:

In[51]:=

Visualize the mean attention across all the attention heads:

In[52]:=

Out[52]=

Visualize each attention head separately:

In[53]:=

Out[53]=

Attention visualization for text

The text feature extractor tokenizes the input string prepending and appending the special tokens StartOfString and EndOfString, adds a special "class" token and then performs causal self-attention on the token embedding vectors. After the self-attention stack, the last vector (corresponding to the "class" token) is used to obtain the final feature representation of the text. Thus the attention procedure for this model can be visualized by inspecting the attention weights between the last vector and the previous ones. Define a test string:

In[54]:=

Extract the NetEncoder of the net to encode the string:

In[55]:=

Out[55]=

In[56]:=

Out[56]=

Extract the list of available tokens and inspect how the input string was tokenized. Even though the BPE tokenization generally segments the input into subwords, it's common to observe that all tokens correspond to full words. Also observe that the StartOfString and EndOfString tokens are added automatically:

In[57]:=

In[58]:=

Out[58]=

In[59]:=

Out[59]=

Feed the string to the net and extract the attention weights used for the last block of self-attention. Note that the original implementation uses a slightly different implementation for the attention mask:

In[60]:=

Out[61]=

Extract the attention weights between the last vector and the previous ones, leaving out the vectors corresponding to StartOfString, EndOfString and the "class" tokens. These weights can be interpreted as which tokens in the original sentence the net is "looking at" in order to perform the feature extraction:

In[62]:=

Out[63]=

Inspect the average attention weights for each token across the attention heads. Observe that the tokens the net is mostly focused on are "hair," "camera" and "picture":

In[64]:=

Out[64]=

Visualize each head separately:

In[65]:=

Out[65]=

Extract the attention weights for all 12 attention layers:

In[66]:=

In[67]:=

Out[68]=

Compute the average across all heads, leaving the StartOfString and EndOfString tokens out:

In[69]:=

Out[70]=

Define a function to visualize the attention weights:

In[71]:=

Explore the attention weights for every layer. A thicker arrow pointing from token A to token B indicates that the layer is paying attention to token B when generating the vector corresponding to token A:

In[72]:=

Out[72]=

Net information

Inspect the number of parameters of all arrays in the net:

In[73]:=

Out[73]=

Obtain the total number of parameters:

In[74]:=

Out[74]=

Obtain the layer type counts:

In[75]:=

Out[75]=