![Revisit consent button](../images/body-df29beacd7496220cc94.svg)
We value your privacy
We use cookies to enhance your browsing experience, serve personalised ads or content, and analyse our traffic. By clicking "Accept All", you consent to our use of cookies.
Customise Reject All Accept All
Customise Consent Preferences ![](../images/body-2159bb67f3de8ad870eb.svg)
We use cookies to help you navigate efficiently and perform certain functions. You will find detailed information about all cookies under each consent category below.
The cookies that are categorised as "Necessary" are stored on your browser as they are essential for enabling the basic functionalities of the site. ... Show more
NecessaryAlways Active
Necessary cookies are required to enable the basic features of this site, such as providing secure log-in or adjusting your consent preferences. These cookies do not store any personally identifiable data.
  * Cookie
BCTempID
  * Duration
10 minutes
  * Description
No description available.


  * Cookie
__cf_bm
  * Duration
1 hour
  * Description
This cookie, set by Cloudflare, is used to support Cloudflare Bot Management. 


  * Cookie
AWSALBCORS
  * Duration
7 days
  * Description
Amazon Web Services set this cookie for load balancing.


  * Cookie
_cfuvid
  * Duration
session
  * Description
Cloudflare sets this cookie to track users across sessions to optimize user experience by maintaining session consistency and providing personalized services


  * Cookie
li_gc
  * Duration
6 months
  * Description
Linkedin set this cookie for storing visitor's consent regarding using cookies for non-essential purposes.


  * Cookie
__hssrc
  * Duration
session
  * Description
This cookie is set by Hubspot whenever it changes the session cookie. The __hssrc cookie set to 1 indicates that the user has restarted the browser, and if the cookie does not exist, it is assumed to be a new session.


  * Cookie
__hssc
  * Duration
1 hour
  * Description
HubSpot sets this cookie to keep track of sessions and to determine if HubSpot should increment the session number and timestamps in the __hstc cookie.


  * Cookie
wpEmojiSettingsSupports
  * Duration
session
  * Description
WordPress sets this cookie when a user interacts with emojis on a WordPress site. It helps determine if the user's browser can display emojis properly.


  * Cookie
BCSessionID
  * Duration
1 year 1 month 4 days
  * Description
Blueconic sets this cookie as a unique identifier for the BlueConic profile.


  * Cookie
_octo
  * Duration
1 year
  * Description
No description available.


  * Cookie
logged_in
  * Duration
1 year
  * Description
No description available.


  * Cookie
__Secure-YEC
  * Duration
past
  * Description
YouTube sets this cookie to stores the user's video player preferences using embedded YouTube video


  * Cookie
__eoi
  * Duration
6 months
  * Description
Description is currently not available.


  * Cookie
AWSALBTGCORS
  * Duration
7 days
  * Description
No description available.


  * Cookie
login-status-p
  * Duration
past
  * Description
Description is currently not available.


  * Cookie
AWSALBTG
  * Duration
7 days
  * Description
No description available.


  * Cookie
csrf_token
  * Duration
session
  * Description
No description available.


  * Cookie
token_v2
  * Duration
1 day
  * Description
Description is currently not available.


  * Cookie
D
  * Duration
1 year
  * Description
Description is currently not available.


  * Cookie
PHPSESSID
  * Duration
session
  * Description
This cookie is native to PHP applications. The cookie stores and identifies a user's unique session ID to manage user sessions on the website. The cookie is a session cookie and will be deleted when all the browser windows are closed.


  * Cookie
VISITOR_PRIVACY_METADATA
  * Duration
6 months
  * Description
YouTube sets this cookie to store the user's cookie consent state for the current domain. 


  * Cookie
cookietest
  * Duration
session
  * Description
The cookietest cookie is typically used to determine whether the user's browser accepts cookies, essential for website functionality and user experience.


  * Cookie
__Host-airtable-session
  * Duration
1 year
  * Description
This cookie is used to enable us to integrate the services of Airtable.


  * Cookie
__Host-airtable-session.sig
  * Duration
1 year
  * Description
This cookie is used to enable us to integrate the services of Airtable.


  * Cookie
m
  * Duration
1 year 1 month 4 days
  * Description
Stripe sets this cookie for fraud prevention purposes. It identifies the device used to access the website, allowing the website to be formatted accordingly.


  * Cookie
BIGipServer*
  * Duration
session
  * Description
Marketo sets this cookie to collect information about the user's online activity and build a profile about their interests to provide advertisements relevant to the user. 


  * Cookie
__cfruid
  * Duration
session
  * Description
Cloudflare sets this cookie to identify trusted web traffic.


  * Cookie
_GRECAPTCHA
  * Duration
6 months
  * Description
Google Recaptcha service sets this cookie to identify bots to protect the website against malicious spam attacks.


  * Cookie
__Secure-YNID
  * Duration
6 months
  * Description
Google cookie used to protect user security and prevent fraud, especially during the login process.


  * Cookie
cookieyes-consent
  * Duration
1 year
  * Description
CookieYes sets this cookie to remember users' consent preferences so that their preferences are respected on subsequent visits to this site. It does not collect or store any personal information about the site visitors.


Functional
Functional cookies help perform certain functionalities like sharing the content of the website on social media platforms, collecting feedback, and other third-party features.
  * Cookie
lidc
  * Duration
1 day
  * Description
LinkedIn sets the lidc cookie to facilitate data center selection.


  * Cookie
brw
  * Duration
1 year
  * Description
No description available.


  * Cookie
brwConsent
  * Duration
5 minutes
  * Description
Description is currently not available.


  * Cookie
WMF-Uniq
  * Duration
1 year
  * Description
Description is currently not available.


  * Cookie
loom_anon_comment
  * Duration
1 year
  * Description
No description available.


  * Cookie
loom_referral_video
  * Duration
session
  * Description
Description is currently not available.


  * Cookie
VISITOR_INFO1_LIVE
  * Duration
6 months
  * Description
A cookie set by YouTube to measure bandwidth that determines whether the user gets the new or old player interface.


  * Cookie
yt-remote-connected-devices
  * Duration
Never Expires
  * Description
YouTube sets this cookie to store the user's video preferences using embedded YouTube videos.


  * Cookie
ytidb::LAST_RESULT_ENTRY_KEY
  * Duration
Never Expires
  * Description
The cookie ytidb::LAST_RESULT_ENTRY_KEY is used by YouTube to store the last search result entry that was clicked by the user. This information is used to improve the user experience by providing more relevant search results in the future.


  * Cookie
yt-remote-device-id
  * Duration
Never Expires
  * Description
YouTube sets this cookie to store the user's video preferences using embedded YouTube videos.


  * Cookie
yt-remote-session-name
  * Duration
session
  * Description
The yt-remote-session-name cookie is used by YouTube to store the user's video player preferences using embedded YouTube video.


  * Cookie
yt-remote-fast-check-period
  * Duration
session
  * Description
The yt-remote-fast-check-period cookie is used by YouTube to store the user's video player preferences for embedded YouTube videos.


  * Cookie
yt-remote-session-app
  * Duration
session
  * Description
The yt-remote-session-app cookie is used by YouTube to store user preferences and information about the interface of the embedded YouTube video player.


  * Cookie
yt-remote-cast-available
  * Duration
session
  * Description
The yt-remote-cast-available cookie is used to store the user's preferences regarding whether casting is available on their YouTube video player.


  * Cookie
yt-remote-cast-installed
  * Duration
session
  * Description
The yt-remote-cast-installed cookie is used to store the user's video player preferences using embedded YouTube video.


  * Cookie
cp_session
  * Duration
3 months
  * Description
Codepen sets this cookie for Help systems found in the website.


  * Cookie
loid
  * Duration
1 year 1 month 4 days
  * Description
This cookie is set by the Reddit. The cookie enables the sharing of content from the website onto the social media platform.


Analytics
Analytical cookies are used to understand how visitors interact with the website. These cookies help provide information on metrics such as the number of visitors, bounce rate, traffic source, etc.
  * Cookie
__hstc
  * Duration
6 months
  * Description
Hubspot set this main cookie for tracking visitors. It contains the domain, initial timestamp (first visit), last timestamp (last visit), current timestamp (this visit), and session number (increments for each subsequent session).


  * Cookie
hubspotutk
  * Duration
6 months
  * Description
HubSpot sets this cookie to keep track of the visitors to the website. This cookie is passed to HubSpot on form submission and used when deduplicating contacts.


  * Cookie
_ga
  * Duration
1 year 1 month 4 days
  * Description
Google Analytics sets this cookie to calculate visitor, session and campaign data and track site usage for the site's analytics report. The cookie stores information anonymously and assigns a randomly generated number to recognise unique visitors.


  * Cookie
_ga_*
  * Duration
1 year 1 month 4 days
  * Description
Google Analytics sets this cookie to store and count page views.


  * Cookie
__Host-psifi.analyticsTrace
  * Duration
6 hours
  * Description
Description is currently not available.


  * Cookie
__Host-psifi.analyticsTraceV2
  * Duration
6 hours
  * Description
Description is currently not available.


  * Cookie
_gh_sess
  * Duration
session
  * Description
GitHub sets this cookie for temporary application and framework state between pages like what step the user is on in a multiple step form.


  * Cookie
YSC
  * Duration
session
  * Description
YSC cookie is set by Youtube and is used to track the views of embedded videos on Youtube pages.


  * Cookie
ajs_anonymous_id
  * Duration
1 year
  * Description
This cookie is set by Segment to count the number of people who visit a certain site by tracking if they have visited before.


  * Cookie
vuid
  * Duration
1 year 1 month 4 days
  * Description
Vimeo installs this cookie to collect tracking information by setting a unique ID to embed videos on the website. 


Performance
Performance cookies are used to understand and analyse the key performance indexes of the website which helps in delivering a better user experience for the visitors.
  * Cookie
AWSALB
  * Duration
7 days
  * Description
AWSALB is an application load balancer cookie set by Amazon Web Services to map the session to the target.


  * Cookie
acq
  * Duration
past
  * Description
Description is currently not available.


  * Cookie
acq.sig
  * Duration
past
  * Description
Description is currently not available.


  * Cookie
ptc
  * Duration
2 years
  * Description
No description available.


Advertisement
Advertisement cookies are used to provide visitors with customised advertisements based on the pages you visited previously and to analyse the effectiveness of the ad campaigns.
  * Cookie
muc_ads
  * Duration
1 year 1 month 4 days
  * Description
Twitter sets this cookie to collect user behaviour and interaction data to optimize the website.


  * Cookie
guest_id_marketing
  * Duration
1 year 1 month 4 days
  * Description
Twitter sets this cookie to identify and track the website visitor.


  * Cookie
guest_id_ads
  * Duration
1 year 1 month 4 days
  * Description
Twitter sets this cookie to identify and track the website visitor.


  * Cookie
personalization_id
  * Duration
1 year 1 month 4 days
  * Description
Twitter sets this cookie to integrate and share features for social media and also store information about how the user uses the website, for tracking and targeting.


  * Cookie
guest_id
  * Duration
1 year 1 month 4 days
  * Description
Twitter sets this cookie to identify and track the website visitor. It registers if a user is signed in to the Twitter platform and collects information about ad preferences.


  * Cookie
bcookie
  * Duration
1 year
  * Description
LinkedIn sets this cookie from LinkedIn share buttons and ad tags to recognize browser IDs.


  * Cookie
__Secure-ROLLOUT_TOKEN
  * Duration
6 months
  * Description
YouTube sets this cookie to manage feature rollout and experimentation. It helps Google control which new features or interface changes are shown to users as part of testing and staged rollouts, ensuring consistent experience for a given user during an experiment.


  * Cookie
yt.innertube::nextId
  * Duration
Never Expires
  * Description
YouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.


  * Cookie
yt.innertube::requests
  * Duration
Never Expires
  * Description
YouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.


  * Cookie
session_tracker
  * Duration
session
  * Description
This cookie is set by the Reddit. This cookie is used to identify trusted web traffic. It also helps in adverstising on the website.


  * Cookie
edgebucket
  * Duration
session
  * Description
Reddit sets this cookie to save the information about a log-on Reddit user, for the purpose of advertisement recommendations and updating the content.


  * Cookie
did
  * Duration
1 year
  * Description
Arbor sets this cookie to show targeted ads to site visitors.This cookie expires after 2 months or 1 year.


Uncategorised
Other uncategorised cookies are those that are being analysed and have not been classified into a category as yet.
No cookies to display.
Reject All Save My Preferences Accept All
Toggle Search[](https://towardsdatascience.com/)Toggle Mobile Navigation
[](https://towardsdatascience.com/)
Publish AI, ML & data-science insights to a global community of data professionals.
Sign in[Submit an Article](https://write.insightmediagroup.io)
[Latest](https://towardsdatascience.com/latest/)[Editor's Picks](https://towardsdatascience.com/tag/editors-pick/)[Newsletter](https://towardsdatascience.com/tag/the-variable/)[ShipAI](https://towardsdatascience.com/shipai/)
* * *
[Write For TDS](https://towardsdatascience.com/submissions/)
[LinkedIn](https://www.linkedin.com/company/towards-data-science/)[X](https://x.com/TDataScience)Toggle Search
[Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/)
# CPTR Paper Walkthrough: Image Captioning, Transformer Mode On
Implementing CPTR (CaPtion TransformeR) from scratch with PyTorch
[Muhammad Ardi Putra](https://towardsdatascience.com/author/muhammad_ardi/)
March 7, 202530 min read
![Black and white photo of a person standing, looking at a wall of art](../images/000.jpg)
Photo by [Antenna](https://unsplash.com/@antenna%20) on [Unsplash](https://unsplash.com/)
## Introduction
In my previous article, I discussed one of the earliest deep learning approaches for image captioning. If you’re interested in reading it, you can find the link to that article at the end of this one.
Today, I would like to talk about image captioning again, but this time with the more advanced neural network architecture. The deep learning I am going to talk about is the one proposed in the paper titled “ _CPTR: Full Transformer Network for Image Captioning_ ,” written by Liu _et al._ back in 2021 [1]. Specifically, here I will reproduce the model proposed in the paper and explain the underlying theory behind the architecture. However, keep in mind that I won’t actually demonstrate the training process since I only want to focus on the model architecture.
### The idea behind CPTR
In fact, the main idea of the CPTR architecture is exactly the same as the earlier image captioning model, as both use the encoder-decoder structure. Previously, in the paper titled “ _Show and Tell: A Neural Image Caption Generator_ ” [2], the models used are GoogLeNet (a.k.a. Inception V1) and LSTM for the two components, respectively. The illustration of the model proposed in the _Show and Tell_ paper is shown in the following figure.
Learn this step by step with the interactive [Machine Learning](https://roadmap.sh/machine-learning) roadmap.
![](../images/001.png)
Figure 1. The neural network architecture for image captioning proposed in the Show and Tell paper [2].
Despite having the same encoder-decoder structure, what makes CPTR different from the previous approach is the basis of the encoder and the decoder themselves. In CPTR, we combine the encoder part of the ViT (Vision Transformer) model with the decoder part of the original Transformer model. The use of transformer-based architecture for both components is essentially where the name CPTR comes from: CaPtion TransformeR.
Note that the discussions in this article are going to be highly related to ViT and Transformer, so I highly recommend you read my previous article about these two topics if you’re not yet familiar with them. You can find the links at the end of this article.
Figure 2 shows what the original ViT architecture looks like. Everything inside the green box is the encoder part of the architecture to be adopted as the CPTR encoder.
![](../images/002.png)
Figure 2. The Vision Transformer (ViT) architecture [3].
Next, Figure 3 displays the original Transformer architecture. The components enclosed in the blue box are the layers that we are going to implement in the CPTR decoder.
![](../images/body-a000d3ae44bd0cf84cc1.png)
Figure 3. The original Transformer architecture [4].
If we combine the components inside the green and blue boxes above, we are going to obtain the architecture shown in Figure 4 below. This is exactly what the CPTR model we are going to implement looks like. The idea here is that the ViT Encoder (green) works by encoding the input image into a specific tensor representation which will then be used as the basis of the Transformer Decoder (blue) to generate the corresponding caption.
![](../images/body-2b651976b51610bf75a3.png)
Figure 4. The CPTR architecture [5].
That’s pretty much everything you need to know for now. I’ll explain more about the details as we go through the implementation.
## Module imports & parameter configuration
As always, the first thing we need to do in the code is to import the required modules. In this case, we only import torch and torch.nn since we are about to implement the model from scratch.
python

```
# Codeblock 1
import torch
import torch.nn as nn
```

Next, we are going to initialize some parameters in Codeblock 2. If you have read my previous article about image captioning with GoogLeNet and LSTM, you’ll notice that here, we got a lot more parameters to initialize. In this article, I want to reproduce the CPTR model as closely as possible to the original one, so the parameters mentioned in the paper will be used in this implementation.
python

```
# Codeblock 2
BATCH_SIZE         = 1              #(1)

IMAGE_SIZE         = 384            #(2)
IN_CHANNELS        = 3              #(3)

SEQ_LENGTH         = 30             #(4)
VOCAB_SIZE         = 10000          #(5)

EMBED_DIM          = 768            #(6)
PATCH_SIZE         = 16             #(7)
NUM_PATCHES        = (IMAGE_SIZE//PATCH_SIZE) ** 2  #(8)
NUM_ENCODER_BLOCKS = 12             #(9)
NUM_DECODER_BLOCKS = 4              #(10)
NUM_HEADS          = 12             #(11)
HIDDEN_DIM         = EMBED_DIM * 4  #(12)
DROP_PROB          = 0.1            #(13)
```

The first parameter I want to explain is the `BATCH_SIZE`, which is written at the line marked with `#(1)`. The number assigned to this variable is not quite important in our case since we are not actually going to train this model. This parameter is set to 1 because, by default, PyTorch treats input tensors as a batch of samples. Here I assume that we only have a single sample in a batch. 
Next, remember that in the case of image captioning we are dealing with images and texts simultaneously. This essentially means that we need to set the parameters for the two. It is mentioned in the paper that the model accepts an RGB image of size 384×384 for the encoder input. Hence, we assign the values for `IMAGE_SIZE` and `IN_CHANNELS` variables based on this information (`#(2)` and `#(3)`). On the other hand, the paper does not mention the parameters for the captions. So, here I assume that the length of the caption is no more than 30 words (`#(4)`), with the vocabulary size estimated at 10000 unique words (`#(5)`).
The remaining parameters are related to the model configuration. Here we set the `EMBED_DIM` variable to 768 (`#(6)`). In the encoder side, this number indicates the length of the feature vector that represents each 16×16 image patch (`#(7)`). The same concept also applies to the decoder side, but in that case the feature vector will represent a single word in the caption. Talking more specifically about the `PATCH_SIZE` parameter, we are going to use the value to compute the total number of patches in the input image. Since the image has the size of 384×384, there will be 576 patches in total (`#(8)`).
When it comes to using an encoder-decoder architecture, it is possible to specify the number of encoder and decoder blocks to be used. Using more blocks typically allows the model to perform better in terms of the accuracy, yet in return, it will require more computational power. The authors of this paper decided to stack 12 encoder blocks (`#(9)`) and 4 decoder blocks (`#(10)`). Next, since CPTR is a transformer-based model, it is necessary to specify the number of attention heads within the attention blocks inside the encoders and the decoders, which in this case authors use 12 attention heads (`#(11)`). The value for the `HIDDEN_DIM` parameter is not mentioned anywhere in the paper. However, according to the ViT and the Transformer paper, this parameter is configured to be 4 times larger than `EMBED_DIM` (`#(12)`). The dropout rate is not mentioned in the paper either. Hence, I arbitrarily set `DROP_PROB` to 0.1 (`#(13)`).
## Encoder
As the modules and parameters have been set up, now that we will get into the encoder part of the network. In this section we are going to implement and explain every single component inside the green box in Figure 4 one by one.
### Patch embedding
![](../images/body-6b159a5730767cc33b10.png)
Figure 5. Dividing the input image into patches and converting them into vectors [5].
You can see in Figure 5 above that the first step to be done is dividing the input image into patches. This is essentially done because instead of focusing on local patterns like CNNs, ViT captures global context by learning the relationships between these patches. We can model this process with the `Patcher` class shown in the Codeblock 3 below. For the sake of simplicity, here I also include the process inside the _patch embedding_ block within the same class.
python

```
# Codeblock 3
class Patcher(nn.Module):
   def __init__(self):
       super().__init__()

       #(1)
       self.unfold = nn.Unfold(kernel_size=PATCH_SIZE, stride=PATCH_SIZE)

       #(2)
       self.linear_projection = nn.Linear(in_features=IN_CHANNELS*PATCH_SIZE*PATCH_SIZE,
                                          out_features=EMBED_DIM)
      
   def forward(self, images):
       print(f'images\t\t: {images.size()}')
       images = self.unfold(images)  #(3)
       print(f'after unfold\t: {images.size()}')
      
       images = images.permute(0, 2, 1)  #(4)
       print(f'after permute\t: {images.size()}')
      
       features = self.linear_projection(images)  #(5)
       print(f'after lin proj\t: {features.size()}')
      
       return features
```

The patching itself is done using the `nn.Unfold` layer (`#(1)`). Here we need to set both the `kernel_size` and `stride` parameters to `PATCH_SIZE (16)` so that the resulting patches do not overlap with each other. This layer also automatically flattens these patches once it is applied to the input image. Meanwhile, the `nn.Linear layer` (`#(2)`) is employed to perform linear projection, i.e., the process done by the _patch embedding_ block. By setting the `out_features` parameter to `EMBED_DIM`, this layer will map every single flattened patch into a feature vector of length 768.
The entire process should make more sense once you read the `forward()` method. You can see at line `#(3)` in the same codeblock that the input image is directly processed by the unfold layer. Next, we need to process the resulting tensor with the `permute()` method (`#(4)`) to swap the first and the second axis before feeding it to the `linear_projection` layer (`#(5)`). Additionally, here I also print out the tensor dimension after each layer so that you can better understand the transformation made at each step.
In order to check if our `Patcher` class works properly, we can just pass a dummy tensor through the network. Look at the Codeblock 4 below to see how I do it.
python

```
# Codeblock 4
patcher  = Patcher()

images   = torch.randn(BATCH_SIZE, IN_CHANNELS, IMAGE_SIZE, IMAGE_SIZE)
features = patcher(images)
```

python

```
# Codeblock 4 Output
images         : torch.Size([1, 3, 384, 384])
after unfold   : torch.Size([1, 768, 576])  #(1)
after permute  : torch.Size([1, 576, 768])  #(2)
after lin proj : torch.Size([1, 576, 768])  #(3)
```

The tensor I passed above represents an RGB image of size 384×384. Here we can see that after the unfold operation is performed, the tensor dimension changed to 1×768×576 (`#(1)`), denoting the flattened 3×16×16 patch for each of the 576 patches. Unfortunately, this output shape does not match what we need. Remember that in ViT, we perceive image patches as a sequence, so we need to swap the 1st and 2nd axes because typically, the 1st dimension of a tensor represents the temporal axis, while the 2nd one represents the feature vector of each timestep. As the `permute()` operation is performed, our tensor is now having the dimension of 1×576×768 (`#(2)`). Lastly, we pass this tensor through the linear projection layer, which the resulting tensor shape remains the same since we set the `EMBED_DIM` parameter to the same size (768) (`#(3)`). Despite having the same dimension, the information contained in the final tensor should be richer thanks to the transformation applied by the trainable weights of the linear projection layer.
### Learnable positional embedding
![](../images/body-90652a40b2d322f3209d.png)
Figure 6. Injecting the learnable positional embeddings into the embedded image patches [5].
After the input image has successfully been converted into a sequence of patches, the next thing to do is to inject the so-called _positional embedding_ tensor. This is essentially done because a transformer without positional embedding is permutation-invariant, meaning that it treats the input sequence as if their order does not matter. Interestingly, since an image is not a literal sequence, we should set the positional embedding to be _learnable_ such that it will be able to somewhat reorder the patch sequence that it thinks works best in representing the spatial information. However, keep in mind that the term “reordering” here does not mean that we physically rearrange the sequence. Rather, it does so by adjusting the embedding weights.
The implementation is pretty simple. All we need to do is just to initialize a tensor using `nn.Parameter` which the dimension is set to match with the output from the `Patcher` model, i.e., 576×768. Also, don’t forget to write `requires_grad=True` just to ensure that the tensor is trainable. Look at the Codeblock 5 below for the details.
python

```
# Codeblock 5
class LearnableEmbedding(nn.Module):
   def __init__(self):
       super().__init__()
       self.learnable_embedding = nn.Parameter(torch.randn(size=(NUM_PATCHES, EMBED_DIM)),
                                               requires_grad=True)
      
   def forward(self):
       pos_embed = self.learnable_embedding
       print(f'learnable embedding\t: {pos_embed.size()}')
      
       return pos_embed
```

Now let’s run the following codeblock to see whether our `LearnableEmbedding` class works properly. You can see in the printed output that it successfully created the positional embedding tensor as expected.
python

```
# Codeblock 6
learnable_embedding = LearnableEmbedding()

pos_embed = learnable_embedding()
```

text

```
# Codeblock 6 Output
learnable embedding : torch.Size([576, 768])
```

### The main encoder block
![](../images/body-fac850f893528cfd337d.png)
Figure 7. The main encoder block [5].
The next thing we are going to do is to construct the main encoder block displayed in the Figure 7 above. Here you can see that this block consists of several sub-components, namely _self-attention_ , _layer norm_ , FFN (Feed-Forward Network), and another _layer norm_. The Codeblock 7a below shows how I initialize these layers inside the `__init__()` method of the `EncoderBlock` class.
python

```
# Codeblock 7a
class EncoderBlock(nn.Module):
   def __init__(self):
       super().__init__()
      
       #(1)
       self.self_attention = nn.MultiheadAttention(embed_dim=EMBED_DIM,
                                                   num_heads=NUM_HEADS,
                                                   batch_first=True,  #(2)
                                                   dropout=DROP_PROB)
      
       self.layer_norm_0 = nn.LayerNorm(EMBED_DIM)  #(3)
      
       self.ffn = nn.Sequential(  #(4)
           nn.Linear(in_features=EMBED_DIM, out_features=HIDDEN_DIM),
           nn.GELU(),
           nn.Dropout(p=DROP_PROB),
           nn.Linear(in_features=HIDDEN_DIM, out_features=EMBED_DIM),
       )
      
       self.layer_norm_1 = nn.LayerNorm(EMBED_DIM)  #(5)
```

I’ve previously mentioned that the idea of ViT is to capture the relationships between patches within an image. This process is done by the _multihead attention_ layer I initialize at line `#(1)` in the above codeblock. One thing to keep in mind here is that we need to set the batch_first parameter to `True` (`#(2)`). This is essentially done so that the attention layer will be compatible with our tensor shape, in which the batch dimension (`batch_size`) is at the 0th axis of the tensor. Next, the two layer normalization layers need to be initialized separately, as shown at line `#(3)` and `#(5)`. Lastly, we initialize the FFN block at line `#(4)`, which the layers stacked using `nn.Sequential` follows the structure defined in the following equation.
![](../images/body-a2ec509a63a7a58620f9.png)
Figure 8. The operations done inside the FFN block [1].
As the `__init__()` method is complete, we will now continue with the `forward()` method. Let’s take a look at the Codeblock 7b below.
python

```
# Codeblock 7b
   def forward(self, features):  #(1)
      
       residual = features  #(2)
       print(f'features & residual\t: {residual.size()}')
      
       #(3)
       features, self_attn_weights = self.self_attention(query=features,
                                                         key=features,
                                                         value=features)
       print(f'after self attention\t: {features.size()}')
       print(f"self attn weights\t: {self_attn_weights.shape}")
      
       features = self.layer_norm_0(features + residual)  #(4)
       print(f'after norm\t\t: {features.size()}')
      

       residual = features
       print(f'\nfeatures & residual\t: {residual.size()}')
      
       features = self.ffn(features)  #(5)
       print(f'after ffn\t\t: {features.size()}')
      
       features = self.layer_norm_1(features + residual)
       print(f'after norm\t\t: {features.size()}')
      
       return features
```

Here you can see that the input tensor is named `features (#(1)`). I name it this way because the input of the `EncoderBlock` is the image that has already been processed with `Patcher` and `LearnableEmbedding`, instead of a raw image. Before doing anything, notice in the `encoder` block that there is a branch separated from the main flow which then returns back to the normalization layer. This branch is commonly known as a _residual connection_. To implement this, we need to store the original input tensor to the residual variable as I demonstrate at line `#(2)`. As the input tensor has been copied, now we are ready to process the original input with the multihead attention layer (`#(3)`). Since this is a _self_ -attention (not a _cross_ -attention), the `query`, `key`, and `value` inputs for this layer are all derived from the `features` tensor. Next, the layer normalization operation is then performed at line `#(4)`, which the input for this layer already contains information from the attention block as well as the residual connection. The remaining steps are basically the same as what I just explained, except that here we replace the self-attention block with FFN (`#(5)`).
In the following codeblock, I’ll test the `EncoderBlock` class by passing a dummy tensor of size 1×576×768, simulating an output tensor from the previous operations.
python

```
# Codeblock 8
encoder_block = EncoderBlock()

features = torch.randn(BATCH_SIZE, NUM_PATCHES, EMBED_DIM)
features = encoder_block(features)
```

Below is what the tensor dimension looks like throughout the entire process inside the model.
python

```
# Codeblock 8 Output
features & residual  : torch.Size([1, 576, 768])  #(1)
after self attention : torch.Size([1, 576, 768])
self attn weights    : torch.Size([1, 576, 576])  #(2)
after norm           : torch.Size([1, 576, 768])

features & residual  : torch.Size([1, 576, 768])
after ffn            : torch.Size([1, 576, 768])  #(3)
after norm           : torch.Size([1, 576, 768])  #(4)
```

Here you can see that the final output tensor (`#(4)`) has the same size as the input (`#(1)`), allowing us to stack multiple encoder blocks without having to worry about messing up the tensor dimensions. Not only that, the size of the tensor also appears to be unchanged from the beginning all the way to the last layer. In fact, there are actually lots of transformations performed inside the attention block, but we just can’t see it since the entire process is done internally by the `nn.MultiheadAttention` layer. One of the tensors produced in the layer that we can observe is the attention weight (`#(2)`). This weight matrix, which has the size of 576×576, is responsible for storing information regarding the relationships between one patch and every other patch in the image. Furthermore, changes in tensor dimension actually also happened inside the FFN layer. The feature vector of each patch which has the initial length of 768 changed to 3072 and immediately shrunk back to 768 again (`#(3)`). However, this transformation is not printed since the process is wrapped with `nn.Sequential` back at line #(4) in Codeblock 7a.
### ViT encoder
![](../images/body-23e7000a7a8f06980ada.png)
Figure 9. The entire ViT Encoder in the CPTR architecture [5].
As we have finished implementing all encoder components, now that we will assemble them to construct the actual ViT Encoder. We are going to do it in the `Encoder` class in Codeblock 9.
python

```
# Codeblock 9
class Encoder(nn.Module):
   def __init__(self):
       super().__init__()
       self.patcher = Patcher()  #(1)
       self.learnable_embedding = LearnableEmbedding()  #(2)

       #(3)
       self.encoder_blocks = nn.ModuleList(EncoderBlock() for _ in range(NUM_ENCODER_BLOCKS))
  
   def forward(self, images):  #(4)
       print(f'images\t\t\t: {images.size()}')
      
       features = self.patcher(images)  #(5)
       print(f'after patcher\t\t: {features.size()}')
      
       features = features + self.learnable_embedding()  #(6)
       print(f'after learn embed\t: {features.size()}')
      
       for i, encoder_block in enumerate(self.encoder_blocks):
           features = encoder_block(features)  #(7)
           print(f"after encoder block #{i}\t: {features.shape}")

       return features
```

Inside the` __init__()` method, what we need to do is to initialize all components we created earlier, i.e., `Patcher` (`#(1)`), `LearnableEmbedding` (`#(2)`), and `EncoderBlock` (`#(3)`). In this case, the `EncoderBlock` is initialized inside `nn.ModuleList` since we want to repeat it `NUM_ENCODER_BLOCKS` (12) times. To the `forward()` method, it initially works by accepting raw image as the input (`#(4)`). We then process it with the `patcher` layer (`#(5)`) to divide the image into small patches and transform them with the linear projection operation. The learnable positional embedding tensor is then injected into the resulting output by element-wise addition (`#(6)`). Lastly, we pass it into the 12 encoder blocks sequentially with a simple for loop (`#(7)`).
Now, in Codeblock 10, I am going to pass a dummy image through the entire encoder. Note that since I want to focus on the flow of this Encoder class, I re-run the previous classes we created earlier with the `print()` functions commented out so that the outputs will look neat.
python

```
# Codeblock 10
encoder = Encoder()

images = torch.randn(BATCH_SIZE, IN_CHANNELS, IMAGE_SIZE, IMAGE_SIZE)
features = encoder(images)
```

And below is what the flow of the tensor looks like. Here, we can see that our dummy input image successfully passed through all layers in the network, including the encoder blocks that we repeat 12 times. The resulting output tensor is now context-aware, meaning that it already contains information about the relationships between patches within the image. Therefore, this tensor is now ready to be processed further with the decoder, which will later be discussed in the subsequent section.
python

```
# Codeblock 10 Output
images                  : torch.Size([1, 3, 384, 384])
after patcher           : torch.Size([1, 576, 768])
after learn embed       : torch.Size([1, 576, 768])
after encoder block #0  : torch.Size([1, 576, 768])
after encoder block #1  : torch.Size([1, 576, 768])
after encoder block #2  : torch.Size([1, 576, 768])
after encoder block #3  : torch.Size([1, 576, 768])
after encoder block #4  : torch.Size([1, 576, 768])
after encoder block #5  : torch.Size([1, 576, 768])
after encoder block #6  : torch.Size([1, 576, 768])
after encoder block #7  : torch.Size([1, 576, 768])
after encoder block #8  : torch.Size([1, 576, 768])
after encoder block #9  : torch.Size([1, 576, 768])
after encoder block #10 : torch.Size([1, 576, 768])
after encoder block #11 : torch.Size([1, 576, 768])
```

### ViT encoder (alternative)
I want to show you something before we talk about the decoder. If you think that our approach above is too complicated, it is actually possible for you to use `nn.TransformerEncoderLayer` from PyTorch so that you don’t need to implement the `EncoderBlock` class from scratch. To do so, I am going to reimplement the `Encoder` class, but this time I’ll name it `EncoderTorch`.
python

```
# Codeblock 11
class EncoderTorch(nn.Module):
   def __init__(self):
       super().__init__()
       self.patcher = Patcher()
       self.learnable_embedding = LearnableEmbedding()
      
       #(1)
       encoder_block = nn.TransformerEncoderLayer(d_model=EMBED_DIM,
                                                  nhead=NUM_HEADS,
                                                  dim_feedforward=HIDDEN_DIM,
                                                  dropout=DROP_PROB,
                                                  batch_first=True)
      
       #(2)
       self.encoder_blocks = nn.TransformerEncoder(encoder_layer=encoder_block,
                                                   num_layers=NUM_ENCODER_BLOCKS)
  
   def forward(self, images):
       print(f'images\t\t\t: {images.size()}')
      
       features = self.patcher(images)
       print(f'after patcher\t\t: {features.size()}')
      
       features = features + self.learnable_embedding()
       print(f'after learn embed\t: {features.size()}')
      
       features = self.encoder_blocks(features)  #(3)
       print(f'after encoder blocks\t: {features.size()}')

       return features
```

What we basically do in the above codeblock is that instead of using the EncoderBlock class, here we use `nn.TransformerEncoderLayer` (`#(1)`), which will automatically create a single encoder block based on the parameters we pass to it. To repeat it multiple times, we can just use nn.`TransformerEncoder` and pass a number to the `num_layers` parameter (`#(2)`). With this approach, we don’t necessarily need to write the forward pass in a loop like what we did earlier (`#(3)`).
The testing code in the Codeblock 12 below is exactly the same as the one in Codeblock 10, except that here I use the `EncoderTorch` class. You can also see here that the output is basically the same as the previous one.
python

```
# Codeblock 12
encoder_torch = EncoderTorch()

images = torch.randn(BATCH_SIZE, IN_CHANNELS, IMAGE_SIZE, IMAGE_SIZE)
features = encoder_torch(images)
```

python

```
# Codeblock 12 Output
images               : torch.Size([1, 3, 384, 384])
after patcher        : torch.Size([1, 576, 768])
after learn embed    : torch.Size([1, 576, 768])
after encoder blocks : torch.Size([1, 576, 768])
```

## Decoder
As we have successfully created the encoder part of the CPTR architecture, now that we will talk about the decoder. In this section I am going to implement every single component inside the blue box in Figure 4. Based on the figure, we can see that the decoder accepts two inputs, i.e., the image caption ground truth (the lower part of the blue box) and the sequence of embedded patches produced by the encoder (the arrow coming from the green box). It is important to know that the architecture drawn in Figure 4 is intended to illustrate the training phase, where the entire caption ground truth is fed into the decoder. Later in the inference phase, we only provide a <BOS> (_Beginning of Sentence_) token for the caption input. The decoder will then predict each word sequentially based on the given image and the previously generated words. This process is commonly known as an _autoregressive_ mechanism.
### Sinusoidal positional embedding
![](../images/body-d327aa2f936ca3b4de94.png)
Figure 10. Where the sinusoidal positional embedding component is located in the decoder [5].
If you take a look at the CPTR model, you’ll see that the first step in the decoder is to convert each word into the corresponding feature vector representation using the _word embedding_ block. However, since this step is very easy, we are going to implement it later. Now let’s assume that this word vectorization process is already done, so we can move to the positional embedding part.
As I’ve mentioned earlier, since transformer is permutation-invariant by nature, we need to apply positional embedding to the input sequence. Different from the previous one, here we use the so-called _sinusoidal positional embedding_. We can think of it like a method to label each word vector by assigning numbers obtained from a sinusoidal wave. By doing so, we can expect our model to understand word orders thanks to the information given by the wave patterns.
If you go back to Codeblock 6 Output, you’ll see that the positional embedding tensor in the encoder has the size of `NUM_PATCHES` × `EMBED_DIM` (576×768). What we basically want to do in the decoder is to create a tensor having the size of `SEQ_LENGTH` × `EMBED_DIM` (30×768), which the values are computed based on the equation shown in Figure 11. This tensor is then set to be non-trainable because a sequence of words must maintain a fixed order to preserve its meaning.
![](../images/body-cac59b5e5c3619142e68.png)
Figure 11. The equation for creating sinusoidal positional encoding proposed in the Transformer paper [6].
Here I want to explain the following code quickly because I actually have discussed this more thoroughly in my previous article about Transformer. Generally speaking, what we basically do here is to create the sine and cosine wave using `torch.sin()` (`#(1)`) and `torch.cos()` (`#(2)`). The resulting two tensors are then merged using the code at line `#(3)` and `#(4)`.
python

```
# Codeblock 13
class SinusoidalEmbedding(nn.Module):
   def forward(self):
       pos = torch.arange(SEQ_LENGTH).reshape(SEQ_LENGTH, 1)
       print(f"pos\t\t: {pos.shape}")
      
       i = torch.arange(0, EMBED_DIM, 2)
       denominator = torch.pow(10000, i/EMBED_DIM)
       print(f"denominator\t: {denominator.shape}")
      
       even_pos_embed = torch.sin(pos/denominator)  #(1)
       odd_pos_embed  = torch.cos(pos/denominator)  #(2)
       print(f"even_pos_embed\t: {even_pos_embed.shape}")
      
       stacked = torch.stack([even_pos_embed, odd_pos_embed], dim=2)  #(3)
       print(f"stacked\t\t: {stacked.shape}")

       pos_embed = torch.flatten(stacked, start_dim=1, end_dim=2)  #(4)
       print(f"pos_embed\t: {pos_embed.shape}")
      
       return pos_embed
```

Now we can check if the `SinusoidalEmbedding` class above works properly by running the Codeblock 14 below. As expected earlier, here you can see that the resulting tensor has the size of 30×768. This dimension matches with the tensor obtained by the process done in the _word embedding_ block, allowing them to be summed in an element-wise manner.
python

```
# Codeblock 14
sinusoidal_embedding = SinusoidalEmbedding()
pos_embed = sinusoidal_embedding()
```

python

```
# Codeblock 14 Output
pos            : torch.Size([30, 1])
denominator    : torch.Size([384])
even_pos_embed : torch.Size([30, 384])
stacked        : torch.Size([30, 384, 2])
pos_embed      : torch.Size([30, 768])
```

### Look-ahead mask
![](../images/body-3a02494252b3e0d4a353.png)
Figure 12. A look-ahead mask needs to be applied to the masked-self attention layer [5].
The next thing I am going to talk about in the decoder is the _masked self-attention_ layer highlighted in the above figure. I am not going to code the attention mechanism from scratch. Rather, I’ll only implement the so-called _look-ahead mask_ , which will be useful for the self-attention layer so that it doesn’t attend to the subsequent words in the caption during the training phase.
The way to do it is pretty easy, what we need to do is just to create a triangular matrix which the size is set to match with the attention weight matrix, i.e., `SEQ_LENGTH` × `SEQ_LENGTH` (30×30). Look at the `create_mask()`function below for the details.
python

```
# Codeblock 15
def create_mask(seq_length):
   mask = torch.tril(torch.ones((seq_length, seq_length)))  #(1)
   mask[mask == 0] = -float('inf')  #(2)
   mask[mask == 1] = 0  #(3)
   return mask
```

Even though creating a triangular matrix can simply be done with `torch.tril()` and `torch.ones()` (`#(1)`), but here we need to make a little modification by changing the 0 values to _-inf_ (`#(2)`) and the 1s to 0 (`#(3)`). This is essentially done because the `nn.MultiheadAttention` layer applies the mask by element-wise addition. By assigning _-inf_ to the subsequent words, the attention mechanism will completely ignore them. Again, the internal process inside an attention layer has also been discussed in detail in[ my previous article about transformer](https://contributor.insightmediagroup.io/paper-walkthrough-attention-is-all-you-need-80399cdc59e1).
Now I am going to run the function with `seq_length=7` so that you can see what the mask actually looks like. Later in the complete flow, we need to set the `seq_length` parameter to `SEQ_LENGTH` (30) so that it matches with the actual caption length.
python

```
# Codeblock 16
mask_example = create_mask(seq_length=7)
mask_example
```

python

```
# Codeblock 16 Output
tensor([[0., -inf, -inf, -inf, -inf, -inf, -inf],
       [0., 0., -inf, -inf, -inf, -inf, -inf],
       [0., 0., 0., -inf, -inf, -inf, -inf],
       [0., 0., 0., 0., -inf, -inf, -inf],
       [0., 0., 0., 0., 0., -inf, -inf],
       [0., 0., 0., 0., 0., 0., -inf],
       [0., 0., 0., 0., 0., 0., 0.]])
```

### The main decoder block
![](../images/body-5425d3ce16026d48365c.png)
Figure 13. The main decoder block [5].
We can see in the above figure that the structure of the decoder block is a bit longer than that of the encoder block. It seems like everything is nearly the same, except that the decoder part has a _cross-attention_ mechanism and an additional layer normalization step placed after it. This cross-attention layer can actually be perceived as the bridge between the encoder and the decoder, as it is employed to capture the relationships between each word in the caption and every single patch in the input image. The two arrows coming from the encoder are the _key_ and _value_ inputs for the attention layer, whereas the _query_ is derived from the previous layer in the decoder itself. Look at the Codeblock 17a and 17b below to see the implementation of the entire decoder block.
python

```
# Codeblock 17a
class DecoderBlock(nn.Module):
   def __init__(self):
       super().__init__()
      
       #(1)
       self.self_attention = nn.MultiheadAttention(embed_dim=EMBED_DIM,
                                                   num_heads=NUM_HEADS,
                                                   batch_first=True,
                                                   dropout=DROP_PROB)
       #(2)
       self.layer_norm_0 = nn.LayerNorm(EMBED_DIM)
       #(3)
       self.cross_attention = nn.MultiheadAttention(embed_dim=EMBED_DIM,
                                                    num_heads=NUM_HEADS,
                                                    batch_first=True,
                                                    dropout=DROP_PROB)

       #(4)
       self.layer_norm_1 = nn.LayerNorm(EMBED_DIM)
      
       #(5)      
       self.ffn = nn.Sequential(
           nn.Linear(in_features=EMBED_DIM, out_features=HIDDEN_DIM),
           nn.GELU(),
           nn.Dropout(p=DROP_PROB),
           nn.Linear(in_features=HIDDEN_DIM, out_features=EMBED_DIM),
       )
      
       #(6)
       self.layer_norm_2 = nn.LayerNorm(EMBED_DIM)
```

In the `__init__()` method, we first initialize both self-attention (`#(1)`) and cross-attention (`#(3)`) layers with `nn.MultiheadAttention`. These two layers appear to be exactly the same now, but later you’ll see the difference in the `forward()` method. The three layer normalization operations are initialized separately as shown at line `#(2)`, `#(4)` and `#(6)`, since each of them will contain different normalization parameters. Lastly, the `ffn` layer (`#(5)`) is exactly the same as the one in the encoder, which basically follows the equation back in Figure 8.
Talking about the `forward()` method below, it initially works by accepting three inputs: `features`, `captions`, and `attn_mask`, which each of them denotes the tensor coming from the encoder, the tensor from the decoder itself, and a look-ahead mask, respectively (`#(1)`). The remaining steps are somewhat similar to that of the `EncoderBlock`, except that here we repeat the multihead attention block twice. The first attention mechanism takes captions as the `query`, `key`, and `value` parameters (`#(2)`). This is essentially done because we want the layer to capture the context within the captions tensor itself — hence the name _self-attention_. Here we also need to pass the attn_mask parameter to this layer so that it cannot see the subsequent words during the training phase. The second attention mechanism is different (`#(3)`). Since we want to combine the information from the encoder and the decoder, we need to pass the `captions` tensor as the `query`, whereas the `features` tensor will be passed as the `key` and `value` — hence the name _cross-attention_. A look-ahead mask is not necessary in the cross-attention layer since later in the inference phase the model will be able to see the entire input image at once rather than looking at the patches one by one. As the tensor has been processed by the two attention layers, we will then pass it through the feed forward network (`#(4)`). Lastly, don’t forget to create the residual connections and apply the layer normalization steps after each sub-component.
python

```
# Codeblock 17b
   def forward(self, features, captions, attn_mask):  #(1)
       print(f"attn_mask\t\t: {attn_mask.shape}")
       residual = captions
       print(f"captions & residual\t: {captions.shape}")
      
       #(2)
       captions, self_attn_weights = self.self_attention(query=captions,
                                                         key=captions,
                                                         value=captions,
                                                         attn_mask=attn_mask)
       print(f"after self attention\t: {captions.shape}")
       print(f"self attn weights\t: {self_attn_weights.shape}")
      
       captions = self.layer_norm_0(captions + residual)
       print(f"after norm\t\t: {captions.shape}")
      
      
       print(f"\nfeatures\t\t: {features.shape}")
       residual = captions
       print(f"captions & residual\t: {captions.shape}")
      
       #(3)
       captions, cross_attn_weights = self.cross_attention(query=captions,
                                                           key=features,
                                                           value=features)
       print(f"after cross attention\t: {captions.shape}")
       print(f"cross attn weights\t: {cross_attn_weights.shape}")
      
       captions = self.layer_norm_1(captions + residual)
       print(f"after norm\t\t: {captions.shape}")
      
       residual = captions
       print(f"\ncaptions & residual\t: {captions.shape}")
      
       captions = self.ffn(captions)  #(4)
       print(f"after ffn\t\t: {captions.shape}")
      
       captions = self.layer_norm_2(captions + residual)
       print(f"after norm\t\t: {captions.shape}")
      
       return captions


```

As the `DecoderBlock` class is completed, we can now test it with the following code.
python

```
# Codeblock 18
decoder_block = DecoderBlock()

features = torch.randn(BATCH_SIZE, NUM_PATCHES, EMBED_DIM)  #(1)
captions = torch.randn(BATCH_SIZE, SEQ_LENGTH, EMBED_DIM)   #(2)
look_ahead_mask = create_mask(seq_length=SEQ_LENGTH)  #(3)

captions = decoder_block(features, captions, look_ahead_mask)
```

Here we assume that features is a tensor containing a sequence of patch embeddings produced by the `encoder` (`#(1)`), while captions is a sequence of embedded words (`#(2)`). The `seq_length` parameter of the look-ahead mask is set to `SEQ_LENGTH` (30) to match it to the number of words in the caption (`#(3)`). The tensor dimensions after each step are displayed in the following output.
python

```
# Codeblock 18 Output
attn_mask             : torch.Size([30, 30])
captions & residual   : torch.Size([1, 30, 768])
after self attention  : torch.Size([1, 30, 768])
self attn weights     : torch.Size([1, 30, 30])    #(1)
after norm            : torch.Size([1, 30, 768])

features              : torch.Size([1, 576, 768])
captions & residual   : torch.Size([1, 30, 768])
after cross attention : torch.Size([1, 30, 768])
cross attn weights    : torch.Size([1, 30, 576])   #(2)
after norm            : torch.Size([1, 30, 768])

captions & residual   : torch.Size([1, 30, 768])
after ffn             : torch.Size([1, 30, 768])
after norm            : torch.Size([1, 30, 768])
```

Here we can see that our `DecoderBlock` class works properly as it successfully processed the input tensors all the way to the last layer in the network. Here I want you to take a closer look at the attention weights at lines `#(1)` and `#(2)`. Based on these two lines, we can confirm that our decoder implementation is correct since the attention weight produced by the self-attention layer has the size of 30×30 (`#(1)`), which basically means that this layer really captured the context within the input caption. Meanwhile, the attention weight matrix generated by the cross-attention layer has the size of 30×576 (`#(2)`), indicating that it successfully captured the relationships between the words and the patches. This essentially implies that after cross-attention operation is performed, the resulting captions tensor has been enriched with the information from the image.
### Transformer decoder
![](../images/body-36e9660848f4d4da495e.png)
Figure 14. The entire Transformer Decoder in the CPTR architecture [5].
Now that we have successfully created all components for the entire decoder, what I am going to do next is to put them together into a single class. Look at the Codeblock 19a and 19b below to see how I do that.
python

```
# Codeblock 19a
class Decoder(nn.Module):
   def __init__(self):
       super().__init__()

       #(1)
       self.embedding = nn.Embedding(num_embeddings=VOCAB_SIZE,
                                     embedding_dim=EMBED_DIM)

       #(2)
       self.sinusoidal_embedding = SinusoidalEmbedding()

       #(3)
       self.decoder_blocks = nn.ModuleList(DecoderBlock() for _ in range(NUM_DECODER_BLOCKS))

       #(4)
       self.linear = nn.Linear(in_features=EMBED_DIM,
                               out_features=VOCAB_SIZE)
```

If you compare this `Decoder` class with the `Encoder` class from codeblock 9, you’ll notice that they are somewhat similar in terms of the structure. In the encoder, we convert image patches into vectors using `Patcher`, while in the decoder we convert every single word in the caption into a vector using the `nn.Embedding layer` (`#(1)`), which I haven’t explained earlier. Afterward, we initialize the positional embedding layer, where for the decoder we use the _sinusoidal_ rather than the _trainable_ one (`#(2)`). Next, we stack multiple decoder blocks using `nn.ModuleList` (`#(3)`). The linear layer written at line #(4), which doesn’t exist in the encoder, is necessary to be implemented here since it will be responsible to map each of the embedded words into a vector of length `VOCAB_SIZE` (10000). Later on, this vector will contain the logit of every word in the dictionary, and what we need to do afterward is just to take the index containing the highest value, i.e., the most likely word to be predicted.
The flow of the tensors within the `forward()` method itself is also pretty similar to the one in the `Encoder` class. In the Codeblock 19b below we pass features, captions, and `attn_mask` as the input (`#(1)`). Keep in mind that in this case the captions tensor contains the raw word sequence, so we need to vectorize these words with the embedding layer beforehand (`#(2)`). Next, we inject the sinusoidal positional embedding tensor using the code at line `#(3)` before eventually passing it through the four decoder blocks sequentially (`#(4)`). Finally, we pass the resulting tensor through the last linear layer to obtain the `prediction` logits (`#(5)`).
python

```
# Codeblock 19b
   def forward(self, features, captions, attn_mask):  #(1)
       print(f"features\t\t: {features.shape}")
       print(f"captions\t\t: {captions.shape}")
      
       captions = self.embedding(captions)  #(2)
       print(f"after embedding\t\t: {captions.shape}")
      
       captions = captions + self.sinusoidal_embedding()  #(3)
       print(f"after sin embed\t\t: {captions.shape}")
      
       for i, decoder_block in enumerate(self.decoder_blocks):
           captions = decoder_block(features, captions, attn_mask)  #(4)
           print(f"after decoder block #{i}\t: {captions.shape}")
      
       captions = self.linear(captions)  #(5)
       print(f"after linear\t\t: {captions.shape}")
      
       return captions
```

At this point you might be wondering why we don’t implement the softmax activation function as drawn in the illustration. This is essentially because during the training phase, softmax is typically included within the loss function, whereas in the inference phase, the index of the largest value will remain the same regardless of whether softmax is applied.
Now let’s run the following testing code to check whether there are errors in our implementation. Previously I mentioned that the captions input of the `Decoder` class is a raw word sequence. To simulate this, we can simply create a sequence of random integers ranging between 0 and `VOCAB_SIZE` (10000) with the length of `SEQ_LENGTH` (30) words (`#(1)`).
python

```
# Codeblock 20
decoder = Decoder()

features = torch.randn(BATCH_SIZE, NUM_PATCHES, EMBED_DIM)
captions = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, SEQ_LENGTH))  #(1)

captions = decoder(features, captions, look_ahead_mask)
```

And below is what the resulting output looks like. Here you can see in the last line that the linear layer produced a tensor of size 30×10000, indicating that our decoder model is now capable of predicting the logit scores for each word in the vocabulary across all 30 sequence positions.
python

```
# Codeblock 20 Output
features               : torch.Size([1, 576, 768])
captions               : torch.Size([1, 30])
after embedding        : torch.Size([1, 30, 768])
after sin embed        : torch.Size([1, 30, 768])
after decoder block #0 : torch.Size([1, 30, 768])
after decoder block #1 : torch.Size([1, 30, 768])
after decoder block #2 : torch.Size([1, 30, 768])
after decoder block #3 : torch.Size([1, 30, 768])
after linear           : torch.Size([1, 30, 10000])
```

### Transformer decoder (alternative)
It is actually also possible to make the code simpler by replacing the `DecoderBlock` class with the `nn.TransformerDecoderLayer`, just like what we did in the ViT Encoder. Below is what the code looks like if we use this approach instead.
python

```
# Codeblock 21
class DecoderTorch(nn.Module):
   def __init__(self):
       super().__init__()
       self.embedding = nn.Embedding(num_embeddings=VOCAB_SIZE,
                                     embedding_dim=EMBED_DIM)
      
       self.sinusoidal_embedding = SinusoidalEmbedding()
      
       #(1)
       decoder_block = nn.TransformerDecoderLayer(d_model=EMBED_DIM,
                                                  nhead=NUM_HEADS,
                                                  dim_feedforward=HIDDEN_DIM,
                                                  dropout=DROP_PROB,
                                                  batch_first=True)
      
       #(2)
       self.decoder_blocks = nn.TransformerDecoder(decoder_layer=decoder_block,
                                                   num_layers=NUM_DECODER_BLOCKS)
      
       self.linear = nn.Linear(in_features=EMBED_DIM,
                               out_features=VOCAB_SIZE)
      
   def forward(self, features, captions, tgt_mask):
       print(f"features\t\t: {features.shape}")
       print(f"captions\t\t: {captions.shape}")
      
       captions = self.embedding(captions)
       print(f"after embedding\t\t: {captions.shape}")
      
       captions = captions + self.sinusoidal_embedding()
       print(f"after sin embed\t\t: {captions.shape}")
      
       #(3)
       captions = self.decoder_blocks(tgt=captions,
                                      memory=features,
                                      tgt_mask=tgt_mask)
       print(f"after decoder blocks\t: {captions.shape}")
      
       captions = self.linear(captions)
       print(f"after linear\t\t: {captions.shape}")
      
       return captions
```

The main difference you will see in the `__init__()` method is the use of `nn.TransformerDecoderLayer` and `nn.TransformerDecoder` at line `#(1)` and `#(2)`, where the former is used to initialize a single decoder block, and the latter is for repeating the block multiple times. Next, the `forward()` method is mostly similar to the one in the `Decoder` class, except that the forward propagation on the decoder blocks is automatically repeated four times without needing to be put inside a loop (`#(3)`). One thing that you need to pay attention to in the `decoder_blocks` layer is that the tensor coming from the encoder (features) must be passed as the argument for the `memory` parameter. Meanwhile, the tensor from the decoder itself (captions) has to be passed as the input to the `tgt` parameter.
The testing code for the `DecoderTorch` model below is basically the same as the one written in Codeblock 20. Here you can see that this model also generates the final output tensor of size 30×10000.
python

```
# Codeblock 22
decoder_torch = DecoderTorch()

features = torch.randn(BATCH_SIZE, NUM_PATCHES, EMBED_DIM)
captions = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, SEQ_LENGTH))

captions = decoder_torch(features, captions, look_ahead_mask)
```

python

```
# Codeblock 22 Output
features             : torch.Size([1, 576, 768])
captions             : torch.Size([1, 30])
after embedding      : torch.Size([1, 30, 768])
after sin embed      : torch.Size([1, 30, 768])
after decoder blocks : torch.Size([1, 30, 768])
after linear         : torch.Size([1, 30, 10000])
```

## The entire CPTR model
Finally, it’s time to put the encoder and the decoder part we just created into a single class to actually construct the CPTR architecture. You can see in Codeblock 23 below that the implementation is very simple. All we need to do here is just to initialize the encoder (`#(1)`) and the decoder (`#(2)`) components, then pass the raw images and the corresponding caption ground truths as well as the look-ahead mask to the `forward()` method (#(3)). Additionally, it is also possible for you to replace the `Encoder` and the `Decoder` with `EncoderTorch` and `DecoderTorch`, respectively.
python

```
# Codeblock 23
class EncoderDecoder(nn.Module):
   def __init__(self):
       super().__init__()
       self.encoder = Encoder()  #EncoderTorch()  #(1)
       self.decoder = Decoder()  #DecoderTorch()  #(2)
      
   def forward(self, images, captions, look_ahead_mask):  #(3)
       print(f"images\t\t\t: {images.shape}")
       print(f"captions\t\t: {captions.shape}")
      
       features = self.encoder(images)
       print(f"after encoder\t\t: {features.shape}")
      
       captions = self.decoder(features, captions, look_ahead_mask)
       print(f"after decoder\t\t: {captions.shape}")
      
       return captions
```

We can do the testing by passing dummy tensors through it. See the Codeblock 24 below for the details. In this case, images is basically just a tensor of random numbers having the dimension of 1×3×384×384 (`#(1)`), while captions is a tensor of size 1×30 containing random integers (`#(2)`).
python

```
# Codeblock 24
encoder_decoder = EncoderDecoder()

images = torch.randn(BATCH_SIZE, IN_CHANNELS, IMAGE_SIZE, IMAGE_SIZE)  #(1)
captions = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, SEQ_LENGTH))  #(2)

captions = encoder_decoder(images, captions, look_ahead_mask)
```

Below is what the output looks like. We can see here that our input images and captions successfully went through all layers in the network, which basically means that the CPTR model we created is now ready to actually be trained on image captioning datasets.
python

```
# Codeblock 24 Output
images         : torch.Size([1, 3, 384, 384])
captions       : torch.Size([1, 30])
after encoder  : torch.Size([1, 576, 768])
after decoder  : torch.Size([1, 30, 10000])
```

## Ending
That was pretty much everything about the theory and implementation of the CaPtion TransformeR architecture. Let me know what deep learning architecture I should implement next. Feel free to leave a comment if you spot any mistakes in this article!
_The code used in this article is available in_[ _my GitHub repo_](https://github.com/MuhammadArdiPutra/medium_articles/blob/main/Image%20Captioning%2C%20Transformer%20Mode%20On.ipynb) _. Here’s the link to my previous article about_[ _image captioning_](https://contributor.insightmediagroup.io/show-and-tell-e1a1142456e2/) _,_[_Vision Transformer (ViT)_](https://contributor.insightmediagroup.io/paper-walkthrough-vision-transformer-vit-c5dcf76f1a7a)_, and the original_[ _Transformer_](https://contributor.insightmediagroup.io/paper-walkthrough-attention-is-all-you-need-80399cdc59e1) _._
## References
[1] Wei Liu _et al._ CPTR: Full Transformer Network for Image Captioning. Arxiv.[ https://arxiv.org/pdf/2101.10804](https://arxiv.org/pdf/2101.10804) [Accessed November 16, 2024].
[2] Oriol Vinyals _et al._ Show and Tell: A Neural Image Caption Generator. Arxiv.[ https://arxiv.org/pdf/1411.4555](https://arxiv.org/pdf/1411.4555) [Accessed December 3, 2024].
[3] Image originally created by author based on: Alexey Dosovitskiy _et al._ An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale. Arxiv.[ https://arxiv.org/pdf/2010.11929](https://arxiv.org/pdf/2010.11929) [Accessed December 3, 2024].
[4] Image originally created by author based on [6].
[5] Image originally created by author based on [1].
[6] Ashish Vaswani _et al._ Attention Is All You Need. Arxiv.[ https://arxiv.org/pdf/1706.03762](https://arxiv.org/pdf/1706.03762) [Accessed December 3, 2024].
## Register Now to Continue Reading
### Unlock Towards Data Science with a Free Account
All articles on Towards Data Science are available to registered users. Create a free account today to access this article and thousands more.
Register for freeSign in to your account
* * *
![](../images/body-59e1c17daab1fcb582c5.jpeg)
Written By
## [Muhammad Ardi Putra](https://towardsdatascience.com/author/muhammad_ardi/)
Deep learning storyteller
[See all from Muhammad Ardi Putra](https://towardsdatascience.com/author/muhammad_ardi/)
[Deep Dives](https://towardsdatascience.com/tag/deep-dives/)[Artificial Intelligence](https://towardsdatascience.com/tag/artificial-intelligence/)[Deep Learning](https://towardsdatascience.com/tag/deep-learning/)[Image Captioning](https://towardsdatascience.com/tag/image-captioning/)[Transformer](https://towardsdatascience.com/tag/transformer/)
Share This Article
  * [](https://www.facebook.com/sharer/sharer.php?u=https%3A%2F%2Ftowardsdatascience.com%2Fimage-captioning-transformer-mode-on%2F&title=CPTR%20Paper%20Walkthrough%3A%20Image%20Captioning%2C%20Transformer%20Mode%20On)
  * [](https://www.linkedin.com/shareArticle?mini=true&url=https%3A%2F%2Ftowardsdatascience.com%2Fimage-captioning-transformer-mode-on%2F&title=CPTR%20Paper%20Walkthrough%3A%20Image%20Captioning%2C%20Transformer%20Mode%20On)
  * [](https://x.com/share?url=https%3A%2F%2Ftowardsdatascience.com%2Fimage-captioning-transformer-mode-on%2F&text=CPTR%20Paper%20Walkthrough%3A%20Image%20Captioning%2C%20Transformer%20Mode%20On)


Towards Data Science is a community publication. Submit your insights to reach our global audience and earn through the TDS Author Payment Program.
[Write for TDS](https://towardsdatascience.com/submissions/)
## Related Articles
  * [![](../images/003.jpg)](https://towardsdatascience.com/you-dont-need-a-job-to-get-data-science-experience-26af4fd4c200/)
## [You Don't Need A Job To Get Data Science Experience](https://towardsdatascience.com/you-dont-need-a-job-to-get-data-science-experience-26af4fd4c200/)
[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)
Flipping the coin on its head
[Kurtis Pykes](https://towardsdatascience.com/author/kurtispykes/)
October 4, 20205 min read
  * [![](../images/004.png)](https://towardsdatascience.com/image-data-generators-in-keras-7c5fc6928400/)
## [Image Data Generators in Keras](https://towardsdatascience.com/image-data-generators-in-keras-7c5fc6928400/)
[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)
How to effectively and efficiently use data generators in Keras for Computer Vision applications of Deep Learning
[Manpreet Singh Minhas](https://towardsdatascience.com/author/manpreetsinghminhas/)
April 24, 20194 min read
  * [![](../images/005.jpg)](https://towardsdatascience.com/can-quantum-machine-learning-help-us-solve-climate-change-3500741d583b/)
## [Can Quantum Machine Learning Help Us Solve Climate Change?](https://towardsdatascience.com/can-quantum-machine-learning-help-us-solve-climate-change-3500741d583b/)
[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)
The history and potential use cases of quantum-derived technology
[Daniel García Solla](https://towardsdatascience.com/author/cardstdani/)
July 13, 202219 min read
  * [![](../images/006.jpeg)](https://towardsdatascience.com/multi-label-land-cover-classification-with-deep-learning-d39ce2944a3d/)
## [Multi-label Land Cover Classification with Deep Learning](https://towardsdatascience.com/multi-label-land-cover-classification-with-deep-learning-d39ce2944a3d/)
[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)
A step by step guide on Classifying Multi-label Land cover classification using Deep Neural Networks
[Abdishakur](https://towardsdatascience.com/author/shakasom/)
November 26, 20195 min read
  * [![](../images/007.jpg)](https://towardsdatascience.com/how-to-use-claude-code-to-build-a-minimum-viable-product/)
## [How to Use Claude Code to Build a Minimum Viable Product](https://towardsdatascience.com/how-to-use-claude-code-to-build-a-minimum-viable-product/)
[Agentic AI](https://towardsdatascience.com/category/artificial-intelligence/agentic-ai/)
Learn how to effectively present product ideas by building MVPs with coding agents
[Eivind Kjosbakken](https://towardsdatascience.com/author/oieivind/)
April 8, 20267 min read
  * [![](../images/008.jpeg)](https://towardsdatascience.com/detecting-scams-using-ai-for-real-6c96e712cf09/)
## [Detecting scams using AI, for real](https://towardsdatascience.com/detecting-scams-using-ai-for-real-6c96e712cf09/)
[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)
A year-long journey to develop one of the best fraud prevention systems on the market
[Massimo Belloni](https://towardsdatascience.com/author/massibelloni/)
July 21, 20197 min read
  * [![](../images/009.png)](https://towardsdatascience.com/understanding-cuped-a822523641af/)
## [Understanding CUPED](https://towardsdatascience.com/understanding-cuped-a822523641af/)
[Data Science](https://towardsdatascience.com/category/data-science/)
An in-depth guide to the state-of-the-art variance reduction technique for A/B tests
[Matteo Courthoud](https://towardsdatascience.com/author/matteo.courthoud/)
June 29, 20224 min read
  * [![](../images/010.png)](https://towardsdatascience.com/i-tried-fine-tuning-a-robot-ai-model-on-colab-here-is-what-worked/)
## [I Fine-Tuned OpenVLA-7B with LoRA on Google Colab: What Worked and What 100 Steps Proved](https://towardsdatascience.com/i-tried-fine-tuning-a-robot-ai-model-on-colab-here-is-what-worked/)
[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)
A reproducible 100-step LoRA fine-tuning run for OpenVLA, with dataset checks, Colab setup, training metrics, and W&B evidence.
[Abdullahi Dattijo](https://towardsdatascience.com/author/abdullahi.dattijo/)
July 21, 202612 min read
  * [![](../images/011.png)](https://towardsdatascience.com/multi-headed-self-attention-by-hand-d2ce1ae031db/)
## [Multi-Headed Self Attention - By Hand](https://towardsdatascience.com/multi-headed-self-attention-by-hand-d2ce1ae031db/)
[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)
Hand computing the cornerstone of modern AI.
[Daniel Warfield](https://towardsdatascience.com/author/danielwarfield1/)
July 12, 20245 min read


  * [YouTube](https://www.youtube.com/c/TowardsDataScience)
  * [X](https://x.com/TDataScience)
  * [LinkedIn](https://www.linkedin.com/company/towards-data-science/)
  * [Threads](https://www.threads.net/@towardsdatascience)
  * [Bluesky](https://bsky.app/profile/towardsdatascience.com)


[](https://towardsdatascience.com/)
Your home for data science and AI. The world's leading publication for data science, data analytics, data engineering, machine learning, and artificial intelligence professionals.
Subscribe to Our Newsletter[Write For TDS](https://towardsdatascience.com/submissions/)•[About](https://towardsdatascience.com/about-towards-data-science-d691af11cc2f/)•[Advertise](https://contact.towardsdatascience.com/advertise-with-towards-data-science)•[Privacy Policy](https://towardsdatascience.com/privacy-policy/)•[Terms of Use](https://towardsdatascience.com/website-terms-of-use/)
© Insight Media Group, LLC 2026•[The New Stack](https://thenewstack.io)•[roadmap.sh](https://roadmap.sh)•[Creator Network](https://insightmediagroup.io/creator-network)
