Nguồn gốc của Vision Transformer

Các bạn theo Computer Vision chắc hẳn đã quá quen mới mạng tích chập truyền thống (CNN), ưu và nhược điểm của mạng này chắc hẳn các bạn đã biết. Để đọc thêm về mạng này có thể truy cập vào đây hoặc đọc bài này của tác giả Phạm Văn Chung. Với sự phát triển của mô hình Transformer bên NLP, các tác giả của Google Research đã cho ra mắt bài báo AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE áp dụng Transformers vào thị giác máy tính.

Về mặt khái niệm, kiến trúc Transformer thường được coi là bất kỳ mạng thần kinh nào sử dụng cơ chế Attention) làm lớp học tập chính của nó. Tương tự như cách mạng nơ-ron tích chập (CNN) sử dụng các tích chập làm lớp học tập chính của nó.

Một số thuật ngữ mình sử dụng bài này:

• ViT: Viết tắt của Vision Transformer, nội dung chính ta sẽ khám phá và tìm hiểu trong bài này.
• Patch Image: Miếng ảnh, được tách từ ảnh gốc.
• Một vài khái niệm khác phía dưới.

1. Một số bước chuẩn bị

• Cài đặt Pytorch. Xuyên suốt bài này mình sẽ sử dụng Pytorch 2.0 : pip3 install torch torchvision torchaudio
• Một số hàm bổ trợ mình sử dụng trong bài, bạn có thể tải tại đây .
• Về dataset mình dùng bộ dataset này, bộ dataset này có hơn 5500 ảnh về động vật, chia thành 90 lớp.
  • Setup thiết bị sử dụng: device = "cuda" if torch.cuda.is_available() else "cpu" device
• Mình cũng sẽ sử dụng Wandb để tracking bài này, thiết lập wandb:

• Cài đặt Pytorch Summary: pip install torch-summary

2. Lấy dữ liệu

• Mình sử dụng bộ dữ liệu từ Roboflow nên mình sử dụng trực tiếp code để download về, khá là tiện, bạn cần thay API Roboflow của bạn để download:

from roboflow import Roboflow
rf = Roboflow(api_key="###Fill Your API-Key")
project = rf.workspace("tbitak-bilgem").project("animalclassification-gktyx")
dataset = project.version(1).download("folder")

• Khi có dataset rồi chúng ta cần lấy đường dẫn train_path và test_path để tiện cho công việc tạo Dataset và Dataloader sau này. Để lấy train_path và test_path mình sử dụng:

image_path = Path("/kaggle/working/AnimalClassification-1")
train_path = image_path.joinpath("train")
test_path = image_path.joinpath("test") 

• Hiển thị ảnh từ dataset thử nào Visualize ảnh từ dataset

Vậy là chúng ta có ảnh đầu vào với kích thước là 640x640x3.

• Thử hiển thị các lớp chúng của bộ dataset nào, mình sẽ lấy tên các lớp dựa vào tên ảnh:

  Mọi người có thể thấy các lớp rồi chứ.

3. Xử lí dữ liệu

Sau khi có dữ liệu, ta cần xử lí một chút để dữ liệu phù hợp với định dạng chuẩn của Pytorch.

3.1. Tạo transform cho hình ảnh

Transform cho hình ảnh trong PyTorch là các chức năng để sửa đổi, tăng cường, hoặc chuẩn hóa hình ảnh theo các cách khác nhau. Mục đích của việc chuyển đổi hình ảnh là để cải thiện chất lượng, đa dạng hóa, hoặc thích ứng hình ảnh với các mô hình khác nhau.

IMG_SIZE=224
from torchvision.transforms import transforms
manual_transform=transforms.Compose(
    [transforms.Resize(size=(IMG_SIZE,IMG_SIZE)),
     transforms.ToTensor()]
)
manual_transform

Hừm, tại sao mình lại chọn 224x224 làm kích thước đầu vào nhỉ?

Trong bài báo tại bảng 3, tác giả có để cập đến Trainning resolution is 224, ở đây chúng ta hiểu ảnh đầu vào có thể hiểu kích thước hình ảnh đầu vào có dạng 224x224x3, đó cũng chính là do mình sử dụng 224x224 làm kích thước đầu vào. Ngoài ra nhóm tác giả cũng đề cập tới việc sử dụng batch_size cho dataloader là 4096, nhưng do giới hạn của phần cứng nên trong bài này mình sẽ sử dụng batch_size=32, giúp tránh việc bị tràn VRAM.

3.2. Tạo Datasets và Dataloaders

Bước tiếp theo chúng ta cần tạo Dataset và Dataloader để phù hợp với yêu cầu đầu vào của Pytorch framework. Phần này mình sẽ xử lí nhanh.

import os
from torchvision import datasets,transforms
from torch.utils.data import DataLoader

NUM_WORKERS=os.cpu_count()
def create_dataloader(train_dir:str,test_dir:str,transform:transforms.Compose,batch_size:int,num_workers:int=NUM_WORKERS):
    train_data=datasets.ImageFolder(train_dir,transform=transform)
    test_data=datasets.ImageFolder(test_dir,transform)
    
    train_dataloader=DataLoader(dataset=train_data,num_workers=num_workers,batch_size=batch_size,shuffle=True,pin_memory=True)
    test_dataloader=DataLoader(dataset=test_data,batch_size=batch_size,pin_memory=True,num_workers=num_workers,shuffle=False)
    class_name=train_data.classes
    return train_dataloader,test_dataloader,class_name
   
BATCH_SIZE=32
train_dataloaders,test_dataloader,class_name=create_dataloader(train_dir=train_dir,test_dir=test_dir,transform=manual_transform,batch_size=BATCH_SIZE,num_workers=0)

Đầu vào của hàm:

---

• train_dir, test_dir : đường dẫn đến train và test đã tạo phía trên.
• transform: phép biến đổi cho hình ảnh, chính là transform mình đã tạo ở trên.
• batch_size: kích thước lô đầu vào, tại đây mình dùng 32.

---

Hàm này sẽ trả về train_dataloader,test_dataloader và class_name của dữ liệu.

Để kiểm tra, dataloader có hoạt động đúng không mình thử visualize một hình ảnh: Trông có vẻ Dataloader đã hoạt động như mong đợi. Trong những phần tiếp theo sẽ rất dài, hi vọng bạn có thể đọc hết.

4. ViT: Tổng quan

Trước tiên chúng ta sẽ tìm hiểu mốt số kiến thức.

4.1. Một số khái niệm

---

Cũng giống như các mạng tích chập, ViT là một kiến trúc mạng sâu, và các mạng nơ-ron thì thường bao gồm các lớp gọi là Layers và các khối Block, một kiến trúc (mô hình) : Architecture(Model)

Các lớp sẽ lấy dữ liệu đầu vào, xử lí, biến đổi chúng,.. và trả ra là các dữ liệu đã được biến đổi. Do đó, nếu một lớp duy nhất lấy đầu vào và cho đầu ra thì một khối sẽ bao gồm nhiều lớp như thế, cũng lấy đầu vào và cho đầu ra. Và một kiến trúc, mô hình thì bao gồm nhiều khối như vậy cộng lại.

---

Bạn có thể nhìn trong hình ảnh:

• Những ô vuông màu xanh lá cây là khối(Block), bao gồm nhiều lớp.
• Ô màu đỏ là lớp (Layers)
• Ô màu dương ngoài cùng chính là mô hình, mô hình bao gồm nhiều khối khác nhau.

4.2. Cấu tạo của ViT

Phần này, mình sẽ mô tả về kiến trúc và cấu tạo của ViT, một số công thức toán đằng sau nó.

4.2.1. Khám phá hình 1 từ paper ViT

Chúng ta có thể thấy kiến trúc ViT bao gồm một số phần:

• Patch + Position Embedding (đầu vào) - Chuyển hình ảnh đầu vào thành một chuỗi các patches và thêm số vị trí để chỉ định thứ tự cho các patches.
• Linear projection of flattened patches (Embedded Patches) : Các patches được chuyển thành Embedding, lợi ích của việc sử dụng Embedding thay vì chỉ các giá trị hình ảnh. Embedding là một biểu diễn có thể học được (thường ở dạng vector) của hình ảnh.
• Norm - Đây là viết tắt của "Layer Normalization" hoặc "LayerNorm", một kỹ thuật để chuẩn hóa (giảm overfitting) một mạng thần kinh, chúng ta có thể sử dụng LayerNorm thông qua lớp PyTorch torch.nn.LayerNorm ().
• Multi-Head Attention - Đây là một lớp Multi-Head Self-Attention hoặc viết tắt là "MSA". Bạn có thể tạo một lớp MSA thông qua lớp PyTorch torch.nn.MultiheadAttention().
• MLP (hoặc Multilayer Perceptron) - MLP thường đề cập đến bất kỳ tập hợp các lớp feedforward nào (hoặc trong trường hợp của PyTorch, một tập hợp các lớp với một phương thức). Trong bài báo ViT, các tác giả gọi MLP là "khối MLP" và nó chứa hai lớp torch.nn.Linear() với kích hoạt phi tuyến tính torch.nn.GELU () ở giữa chúng và lớp torch.nn.Dropout().
• Transformer Encoder - Bộ mã hóa Transformer, là một tập hợp các lớp được liệt kê ở trên. Có hai kết nối bỏ qua bên trong bộ mã hóa Transformer (ký hiệu "+") có nghĩa là đầu vào của lớp được đưa trực tiếp đến các lớp ngay lập tức cũng như các lớp tiếp theo. Kiến trúc ViT tổng thể bao gồm một số bộ mã hóa Transformer xếp chồng lên nhau.
• MLP Head - Đây là lớp đầu ra của kiến trúc, nó chuyển đổi các features đã học của đầu vào thành đầu ra lớp. Vì chúng ta đang nghiên cứu phân loại hình ảnh, bạn cũng có thể gọi đây là "đầu phân loại". Cấu trúc của MLP Head tương tự như khối MLP bên trên.

Một số thuật ngữ:

• Patches: Hình ảnh đã được chia nhỏ, hay còn gọi là miếng ảnh
• Embedding: kĩ thuật biểu diễn các Image Patches dưới dạng vector có số chiều nhỏ hơn so với kích thước của image patches nhưng vẫn bảo toàn được các mối quan hệ và ngữ nghĩa của hình ảnh
• Flattened patches là các miếng ảnh đã được làm phẳng và chuyển thành các vector một chiều. Mỗi miếng ảnh có kích thước nhỏ hơn so với ảnh gốc, nhưng vẫn giữ được các thông tin cơ bản về màu sắc, độ sáng, độ tương phản,..

4.2.2. Khám phá 4 công thức

Tiếp theo chúng ta sẽ đi tìm hiểu 4 công thức mà nhóm tác giả đã đề cập trong bài báo.

Mô tả các công thức theo ViT paper.

Công thức	Mô tả từ phần 3.1 của bài báo ViT
1	Transformer sử dụng kích thước vector ẩn D cố định qua tất cả các lớp của nó, vì vậy chúng tôi làm phẳng các mảnh và ánh xạ vào D chiều với một phép chiếu tuyến tính có thể đào tạo (Công thức 1). Đầu ra của phép chiếu này là các patch embeddings... Các Position Embedding được thêm vào các nhúng mảnh để giữ thông tin vị trí. Chúng tôi sử dụng các nhúng vị trí 1D có thể học theo cách tiêu chuẩn...
2	Bộ mã hóa Transformer (Vaswani et al., 2017) bao gồm các lớp xen kẽ của tự chú ý đa đầu (MSA, xem Phụ lục A) và các khối MLP (Công thức 2, 3). Layernorm (LN) được áp dụng trước mỗi khối, và các kết nối dư sau mỗi khối (Wang et al., 2019; Baevski & Auli, 2019).
3	Giống như công thức 2.
4	Tương tự như token [ class ] của BERT, chúng tôi thêm một nhúng có thể học vào chuỗi các mảnh đã nhúng $\left(\mathbf{z}_{0}^{0}=\mathbf{x}_{\text {class }}\right)$, trạng thái của nó tại đầu ra của bộ mã hóa Transformer $\left(\mathbf{z}_{L}^{0}\right)$ phục vụ là biểu diễn hình ảnh $\mathbf{y}$ (Công thức 4)...

---

Chúng ta có thể mapping cho từng công thức với các phần trong mô hình. Mapping từng công thức với các phần trong mô hình

---

Trong tất cả các công thức, trừ công thức 4, "$\mathbf{z}$" là đầu ra của một lớp cụ thể.

1. $\mathbf{z}_{0}$ "z zero" (đây là đầu ra của lớp patch embedding ban đầu)
2. $\mathbf{z}_{\ell}^{\prime}$ là "z của một số nguyên tố lớp cụ thể" (hoặc giá trị trung gian của z).
3. $\mathbf{z}_{\ell}$ "z của một lớp cụ thể". Và $\mathbf{y}$ là đầu ra tổng thể của mô hình.

4.2.3. Tổng quan về công thức 1

Công thức 1:

$$\begin{aligned} \mathbf{z}_{0} &=\left[\mathbf{x}_{\text {class }} ; \mathbf{x}_{p}^{1} \mathbf{E} ; \mathbf{x}_{p}^{2} \mathbf{E} ; \cdots ; \mathbf{x}_{p}^{N} \mathbf{E}\right]+\mathbf{E}_{\text {pos }}, & & \mathbf{E} \in \mathbb{R}^{\left(P^{2} \cdot C\right) \times D}, \mathbf{E}_{\text {pos }} \in \mathbb{R}^{(N+1) \times D} \end{aligned}$$

Công thức này liên quan đến class token, patch embedding và position embedding ( E là Embedding) của hình ảnh đầu vào.

Ở dạng vector, việc nhúng có thể trông giống như: x_input = [class_token, image_patch_1, image_patch_2, image_patch_3...] + [class_token_position, image_patch_1_position, image_patch_2_position, image_patch_3_position...]

4.2.4. Tổng quan về công thức 2

Công thức 2:

$$\begin{aligned} \mathbf{z}_{\ell}^{\prime} &=\operatorname{MSA}\left(\operatorname{LN}\left(\mathbf{z}_{\ell-1}\right)\right)+\mathbf{z}_{\ell-1}, & & \ell=1 \ldots L \end{aligned}$$

Với mọi lớp từ 1 tới L(tổng số lớp), có một lớp Multi-Head Attention (MSA) chứa một lớp LayerNorm (LN). Chúng ta sẽ gọi lớp này là "MSA block".

Biểu diễn trong pseudo code: x_output_MSA_block = MSA_layer(LN_layer(x_input)) + x_input

4.2.5. Tổng quan về công thức 3.

Công thức 3:

$$\begin{aligned} \mathbf{z}_{\ell} &=\operatorname{MLP}\left(\operatorname{LN}\left(\mathbf{z}_{\ell}^{\prime}\right)\right)+\mathbf{z}_{\ell}^{\prime}, & & \ell=1 \ldots L \\ \end{aligned}$$

Tương tự như công thức phía trên, với mọi lớp từ 1 thông qua đến L (tổng số lớp), cũng có một lớp Perceptron đa lớp (MLP) chứa một lớp LayerNorm (LN). Mình sẽ gọi layer này là "MLP block".

Biểu diễn bằng pseudo code: x_output_MLP_block = MLP_layer(LN_layer(x_output_MSA_block)) + x_output_MSA_block

4.2.6. Tổng quan về công thức 4.

Công thức 4:

$$\begin{aligned} \mathbf{y} &=\operatorname{LN}\left(\mathbf{z}_{L}^{0}\right) & & \end{aligned}$$

Công thức này có đầu ra y là index 0 token của z, được chứa trong một lớp LayerNorm(LN). Mình sẽ gọi công thức này là x_output_MLP_block.

Pseudo code : y = Linear_layer(LN_layer(x_output_MLP_block[0]))

4.2.7. Khám phá bảng 1

Phần cuối cùng của mô hình ViT, chúng ta sẽ đi khám phá bảng 1.

Model	Layers	Hidden size $D$	MLP size	Heads	Params
ViT-Base	12	768	3072	12	$86M$
ViT-Large	24	1024	4096	16	$307M$
ViT-Huge	32	1280	5120	16	$632M$

Chúng ta có thể hiểu bảng 1 đơn giản như sau:

• Layers - Số lượng khối Transformer Encoder (mỗi khối trong số này sẽ chứa một khối MSA và khối MLP)
• Hidden size D: Đây là Embedding Size xuyên suốt mô hình. Đơn giản hơn, đây sẽ là kích thước của vector mà hình ảnh của chúng ta được biến thành khi nó được patched và embedding. Nói chung, kích thước embedding càng lớn, càng có thể nắm bắt được nhiều thông tin, kết quả càng tốt. Tuy nhiên, việc nhúng lớn hơn đi kèm với khối lượng tính toán lớn hơn.
• MLP Size - Số lượng đơn vị ẩn trong các lớp MLP.
• Heads - Số lượng đầu trong các layer Multi-Head Attention
• Params - Tổng số tham số của mô hình. Nói một cách đơn giản, nhiều tham số hơn dẫn đến hiệu suất tốt hơn nhưng khối lượng tính toán cũng sẽ lớn hơn.

5. Bắt đầu code thôi nào

5.1.1. Công thức 1: Tách ảnh ban đầu thành các miếng (patches), tạo các lớp, vị trí nhúng,...

---

Mở đầu phần 3.1 của bài báo AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE nhóm tác giả có viết như sau

Chúng ta có thể hiểu đoạn này một các đơn giản:

• D là kích thước của các Patch Embedding, các giá trị của D có thể xem tại phần 4.2.7 mình đã nói ở trên

• Hình ảnh đầu vào sẽ bắt đầu dưới dạng 2D với kích thước ${H \times W \times C}$. Trong đó

  • $(H, W)$ là độ phân giải của ảnh gốc .
  • $C$ số lượng kênh.

• Hình ảnh được chuyển đổi thành một chuỗi Flatten 2D Patches với kích thước ${N \times\left(P^{2} \cdot C\right)}$.

  • $(P, P)$ là độ phân giải cho từng miếng ảnh (patch size)
  • $N=H W / P^{2}$ là số lượng patch đầu ra, cũng chính là đầu vào cho mô hình Transformer.
  

Mapping patch và position embedding với mô hình ViT . Source: Mrdbourke

5.1.2. Thử tính toán đầu vào và đầu ra của Patch Embedding

Trước tiên chúng ta hãy cài đặt đầu vào cho mô hình:

height=224 #Chiều cao của ảnh
width=224 #Chiều rộng của ảnh
color_channels=3 #Kênh màu
patch_size=16 #Patch size

number_of_patches=int((height*width)/patch_size**2)
print(f"Number of patches (N) with image height (H={height}), width (W={width}) and patch size (P={patch_size}): {number_of_patches}")

Output: Number of patches (N) with image height (H=224), width (W=224) and patch size (P=16): 196

Chúng ta có thể tính patches của ảnh bằng cách dùng công thức $N=H W / P^{2}$ hay $N=(224*224)/16^{2}=196$

Chúng ta đã có số lượng miếng ảnh, làm thế nào để chúng ta tìm được kích thước một chuỗi các flatten 2D patches. Chúng ta có công thức sau

• Đầu vào: Hình ảnh bắt đầu với 2D có kích thước ${H \times W \times C}$.
• Đầu ra: Hình ảnh chuyển đổi thành flatten 2D patches với kích thước ${N \times\left(P^{2} \cdot C\right)}$.

Chúng ta có thể dựa vào công thức phía trên để viết code:

embedding_layer_input_shape=(height,width,color_channels)
embedding_layer_output_shape=(number_of_patches,patch_size**2*color_channels)
print(f"Input shape (single 2D image): {embedding_layer_input_shape}")
print(f"Output shape (single 2D image flattened into patches): {embedding_layer_output_shape}")

Output: 
Input shape (single 2D image): (224, 224, 3)
Output shape (single 2D image flattened into patches): (196, 768)

5.1.3. Biến hình ảnh các patches

Chúng ta có có ảnh đầu vào như sau:

Làm thế nào để biến hình ảnh này thành các patches với hình 1 của ViT paper. Chúng ta có thể làm điều này bằng cách lập chỉ mục trên các kích thước hình ảnh khác nhau.

Chúng ta đã có hàng trên cùng, hãy biến thành thành các patch:

Chúng ta đã biến hàng trên cùng thành các patch, vậy còn toàn bộ ảnh thì sao. Hãy thử đoạn code sau:

img_size = 224
patch_size = 16
num_patches = img_size/patch_size
assert img_size % patch_size == 0, "Image size must be divisible by patch size"
print(f"Number of patches per row: {num_patches}\
        \nNumber of patches per column: {num_patches}\
        \nTotal patches: {num_patches*num_patches}\
        \nPatch size: {patch_size} pixels x {patch_size} pixels")
fig, axs = plt.subplots(nrows=img_size // patch_size,
                        ncols=img_size // patch_size,
                        figsize=(num_patches, num_patches),
                        sharex=True,
                        sharey=True)

for i, patch_height in enumerate(range(0, img_size, patch_size)): 
    for j, patch_width in enumerate(range(0, img_size, patch_size)):
        axs[i, j].imshow(image_permuted[patch_height:patch_height+patch_size, 
                                        patch_width:patch_width+patch_size, 
                                        :]) 
        axs[i, j].set_ylabel(i+1,
                             rotation="horizontal",
                             horizontalalignment="right",
                             verticalalignment="center")
        axs[i, j].set_xlabel(j+1)
        axs[i, j].set_xticks([])
        axs[i, j].set_yticks([])
        axs[i, j].label_outer()

# Set a super title
fig.suptitle(f"{class_name[label]} -> Patchified", fontsize=16)
plt.show()

Đầu ra sẽ là ảnh được biến thành các patch như dưới đây:

Hừm, đã có các patch, làm thế nào để biến chúng thành các Embedding và một chuỗi. Chúng ta có thể dùng các lớp có sẵn của Pytorch.

5.1.4. Tạo các patches image

Ở phần trên chúng ta đã thấy hình ảnh được biến thành các patches, vậy chúng ta cần phải chuyển chúng trong Pytorch. Đọc paper, tại phần 3.1, tác giả có đề cập tới kiến trúc lai

Vậy là chúng ta có thể CNN để tạo ra các patch image bằng cách thiết lập `kernel_size` và `stride` của `torch.nn.Conv2d()` bằng với `patch size`. Bạn có thể xem GIF dưới đây để dễ hình dung:

Biến ảnh thành patches . Source: Mrdbourke

Đã có ý tưởng, chúng ta có thể code:

from torch import nn
patch_size=16
conv2d=nn.Conv2d(in_channels=3,out_channels=768,kernel_size=patch_size,stride=patch_size,padding=0)

• Đầu vào là kênh màu, tại đây là 3
• Đầu ra là 768, theo như bảng 1, xem lại phần 4.2.7. Đây là kích thước Embedding, mỗi hình ảnh sẽ được nhúng vào một vectơ có thể học được có kích thước 768

Bây giờ chúng ta hãy kiểm tra kích thước của ảnh sau khi dùng conv2d Mình đã dùng unsqeeze để thêm chiều batch_size.

Vậy là chúng ta đã có các patches image. Tiếp theo cần làm flatten cho patch embedding. Đầu ra của chúng ta có kích thước lần lượt là torch.Size([1, 768, 14, 14]) -> [batch_size, embedding_dim, feature_map_height, feature_map_width]. Ảnh đã được chuyển thành các patches_image với kích thước 14x14.

Thử hiển thị Feature Maps nào:

Các feature maps này có thể thay đổi khi mạng thần kinh học, Do đó nó có thể được coi là Embedding có thể học được.

5.1.5. Flatten cho các patches image

Chúng ta đã biến hình ảnh thành các Patch Embedding nhưng chúng vẫn đang ở 2D.Trong Pytorch có một hàm chúng ta có thể làm phẳng cho image là torch.nn.Flatten(). Đầu ra mà chúng ta mong muốn là : (196, 768) -> (number of patches, embedding dimension) hay ${N \times\left(P^{2} \cdot C\right)}$

Trong paper ViT, (đọc lại ảnh phần đầu 5.1.4), nhóm tác giả có 1 đoạn viết:

As a special case, the patches can have spatial size 1x1, which means that the input sequence is obtained by simply flattening the spatial dimensions of the feature map and projecting to the Transformer dimension. The classification input embedding and position embeddings are added as described above.

Chúng ta sẽ không làm phẳng toàn bộ tensor mà chỉ làm phẳng phần đặc biệt chính là spatial dimensions of the feature map( Phần này mình không mô tả bằng tiếng Việt được ). Trong trường hợp này là kích thước kích thuóc của feature_map_height feature_map_width image_out_of_conv. Chúng ta có thể dùng start_dim và end_dim để flatten cục bộ.

flatten=nn.Flatten(start_dim=2 # flatten feature_map_height (dimension 2)
                   ,end_dim=3) #flatten feature map_width(dimension 3)

Chúng ta có thể ghép chúng lại với nhau:

Có vẻ hình dạng của chúng ta bị ngược với hình dạng mà chúng ta mong muốn (196, 768) nhỉ? Chúng ta có thể sắp xếp lại số chiều, sử dụng permute

image_out_of_conv_flattened_reshaped = image_out_of_conv_flattend.permute(0, 2, 1) # [batch_size, P^2•C, N] -> [batch_size, N, P^2•C]
print(f"Patch embedding sequence shape: {image_out_of_conv_flattened_reshaped.shape} -> [batch_size, num_patches, embedding_size]")

Output: Patch embedding sequence shape: torch.Size([1, 196, 768]) -> [batch_size, num_patches, embedding_size]

Hãy hiển thị kết quả của chúng ta nào:

Vậy là code đã hoạt động.

5.1.6. Chuyển ViT Patch Embedding vào module Pytorch.

Vậy là chúng ta đã trải qua một hành trình khá dài, mình cùng chuyển tất cả vào một module để tiện thao tác duy nhất thôi. Chúng ta sẽ:

1. Tạo một lớp cha chứa các lớp con (vì vậy nó có thể được sử dụng một lớp PyTorch).PatchEmbeddingnn.Module
2. Khởi tạo lớp với các tham số , (đối với ViT-Base) và (đây là in_channels=3 patch_size=16 embedding_dim=768 D cho ViT-Base từ Bảng 1).
3. Tạo một lớp để biến hình ảnh thành các patch bằng cách sử dụng nn.Conv2d()(giống như trong 5.1.5 ở trên)
4. Tạo một lớp để làm phẳng các patch feature maps thành một chiều duy nhất (giống như trong 4.4 ở trên). Xác định một phương thức để lấy một đầu vào và truyền nó qua các lớp được tạo trong 3 và 4. forward()
5. Đảm bảo hình dạng đầu ra phản ánh hình dạng đầu ra cần thiết của kiến trúc ViT ${N \times\left(P^{2} \cdot C\right)}$.

Code thôi nào:

# 1. Create a class which subclasses nn.Module
class PatchEmbedding(nn.Module):
    '''Turn a 2D input into 1D sequence learnable embedding vector
    Args:
        - in_channels(int): Number of color channels for the input images. Defaults to 3
        - Patch_size (int): Size of patches to convert input image into. Defaults to 16
        - embedding_dim(int): Size of embedding to turn image. Default to 768
    '''
    #2. Initalize the class with apporpriate variables
    def __init__(self,in_channel:int=3,patch_size:int=16,embedding_dim:int=768):
        super().__init__()
        #3. Create a layer to turn an image into patches
        self.patcher=nn.Conv2d(in_channels=in_channel,padding=0,stride=patch_size,out_channels=embedding_dim,kernel_size=patch_size)
        #4. Create a layer to flatten the patch features maps into a single dimension
        self.flatten=nn.Flatten(start_dim=2,end_dim=3)
    def forward(self,x):
        image_resolution=x.shape[-1]
        assert image_resolution%patch_size==0, f"Input image size must be divisible by patch size, image shape:{image_resolution}, patch size: {patch_size}"
        
        #Perform the forward pass
        x_patched=self.patcher(x)
        x_flattend=self.flatten(x_patched)
        return x_flattend.permute(0,2,1) # Adjust so the embedding is on the final dimension [batch_size, P^2•C, N] -> [batch_size, N, P^2•C]

Vậy là đã tạo xong lớp PatchEmbedding, hãy kiểm tra đoạn code nào:

def set_seed(seed:int=42):
  torch.manual_seed(seed)
  torch.cuda.manual_seed(seed)

set_seed()
patchify=PatchEmbedding(in_channel=3,patch_size=16,embedding_dim=768)
print(f"Input image shape: {image.unsqueeze(0).shape}")
patch_embedded_image=patchify(image.unsqueeze(0)) #Add an extra batch dimension
print(f"Output patch embedding shape: {patch_embedded_image.shape}")

Output: Input image shape: torch.Size([1, 3, 224, 224])
Output patch embedding shape: torch.Size([1, 196, 768])

Tổng quan về tham số của lớp:

Vậy là đầu ra đã như những gì chúng ta mong muốn. Điểm lại những gì đã làm với mô hình nào:

Lớp PatchEmbedding của chúng ta (bên phải) tái tạo lại embedding cho các miếng ảnh của kiến trúc ViT từ Hình 1 và Công thức 1 trong bài báo ViT (bên trái). Tuy nhiên, embedding cho lớp có thể học được và embedding cho vị trí chưa được tạo ra. Source: Mrdbourke

5.1.7. Tạo Class Token Embedding

Đoạn 2 phần 3 của paper ViT có viết như sau

Source: Paper ViT

BERT (Bidirectional Encoder Representations from Transformers) là một trong những tài liệu nghiên cứu học máy ban đầu sử dụng kiến trúc Transformer để đạt được kết quả nổi bật về các nhiệm vụ xử lý ngôn ngữ tự nhiên (NLP) và là nơi bắt nguồn ý tưởng có token thông báo ở đầu chuỗi, lớp là mô tả cho lớp "classification" mà chuỗi thuộc về.[ class ]. Để biết thông tin thêm về BERT bạn có xem paper này hoặc bài Viblo này

Vì vậy, chúng ta cần phải "chuẩn bị Embedding có thể học được vào chuỗi các Patch Embedding".

Chúng ta hãy xem lại chuỗi các Embedding ( tạo trong 5.1.6) và kích thước của nó

Để "chuẩn bị Embedding có thể học được vào chuỗi các Patch Embedding", chúng ta cần tạo ra một Embedding có thể học được trong shape của embedding_dimension (D) sau đó cộng chúng vào chiều number_of_patches

Trong pseudo code: patch_embedding = [image_patch_1, image_patch_2, image_patch_3...] class_token = learnable_embedding patch_embedding_with_class_token = torch.cat((class_token, patch_embedding), dim=1)

Code của chúng ta như sau:

batch_size=patch_embedded_image.shape[0]
embedding_dimension=patch_embedded_image.shape[-1]

class_token=nn.Parameter(torch.ones(batch_size,1,embedding_dimension),# [batch_size, number_of_tokens, embedding_dimension]
                        requires_grad=True)

print(class_token[:,:,10])
print(f"Shape of classtoken: {class_token.shape}->[batch_size, number_of_tokens, embedding_dimension]")

Output: tensor([[1.]], grad_fn=<SelectBackward0>)
Shape of classtoken: torch.Size([1, 1, 768])->[batch_size, number_of_tokens, embedding_dimension]

Trong đoạn code mình đã dùng torch.ones, mục đích là để có thể biểu diễn, trong thực tế có thể thay đổi thành torch.randn() để khai thác hết súc mạng của random.

Vậy là ta đã tạo Embedding có thể học được, bây giờ cần gộp chúng vào Patch Embedding. Chúng ta có thể dùng torch.cat để gộp chúng.

# Add the class token embedding to the front of the patch embedding
patch_embedded_image_with_class_embedding = torch.cat((class_token, patch_embedded_image),
                                                      dim=1) # concat on first dimension

# Print the sequence of patch embeddings with the prepended class token embedding
print(patch_embedded_image_with_class_embedding)
print(f"Sequence of patch embeddings with class token prepended shape: {patch_embedded_image_with_class_embedding.shape} -> [batch_size, number_of_patches, embedding_dimension]")

Output

Sequence of patch embeddings with class token prepended shape: > torch.Size([1, 197, 768]) -> [batch_size, number_of_patches, embedding_dimension]

Bạn có thể nhìn ảnh phía dưới để dễ hình dung về quá trình thao tác:

Source: Mrdbourke

---

5.1.8. Tạo Position Embedding

Trước khi chúng ta bắt đầu bạn có thể nhìn ảnh sau:

$\mathbf{E}_{\text {pos }}$ từ công thức 1 $E$ viết tắt của "embedding".

Source: Mrdbourke

Nhóm tác giả có viết, trích đoạn 3 phần 3.1

Position embeddings are added to the patch embeddings to retain positional information. We use standard learnable 1D position embeddings, since we have not observed significant performance gains from using more advanced 2D-aware position embeddings (Appendix D.4). The resulting sequence of embedding vectors serves as input to the encoder.

Bằng cách “giữ thông tin Positional” các tác giả muốn kiến trúc biết được “thứ tự” của các miếng ảnh. Nghĩa là, miếng ảnh thứ hai đến sau miếng ảnh thứ nhất và miếng ảnh thứ ba đến sau miếng ảnh thứ hai và cứ tiếp tục như vậy.

Thông tin vị trí này có thể quan trọng khi xem xét những gì có trong một bức ảnh (nếu không có thông tin vị trí, một chuỗi phẳng có thể được coi như không có thứ tự và do đó không có miếng ảnh nào liên quan đến miếng ảnh nào khác).

Trước khi bắt đầu tạo Position Embedding, hãy xem Embedding mà chúng ta đang có:

Công thức 1 cũng nói răng Positional Embedding ($\mathbf{E}_{\text {pos }}$) nên có kích thước $(N + 1) \times D$:

$$\mathbf{E}_{\text {pos }} \in \mathbb{R}^{(N+1) \times D}$$

Trong đó:

• $N=H W / P^{2}$ là số lượng miếng ảnh(Patches) kết quả, cũng là độ dài chuỗi đầu vào hiệu quả cho Transformer (số lượng Patches)..
• $D$là kích thước của các patch embeddings, các giá trị khác nhau cho $D$ ó thể được tìm thấy trong Bảng 1(embedding dimension).

Chúng ta có thể tạo Position Embedding với đoạn code dưới đây:

# Calculate N ( number of patches)
number_of_patches=int((height*width)/patch_size**2)
#Get Embedding dimension
embedding_dimension=patch_embedded_image_with_class_embedding.shape[2]
#Create the learneable 1D position embedding
position_embedding=nn.Parameter(torch.ones(1,number_of_patches+1,embedding_dimension),requires_grad=True)
print(position_embedding[:,:,10])
print(f"Position embedding shape: {position_embedding.shape}->[batch_size, number_of_patches, embedding_dimension]")

Output: Position embeddding shape: torch.Size([1, 197, 768]) -> [batch_size, number_of_patches, embedding_dimension]

Positional Embedding đã tạo xong, chúng ta chỉ cân thêm vào Class Token Embedding phía trên:

Bạn có thể hình dung những gì ta đã làm bằng cách xem ảnh phía dưới:

Source: Mrdbourke

5.1.9. Kết hợp lại tất cả với nhau

Chúng ta đã tìm hiểu và khám phá xong công thức 1, đã đến lúc biến chúng thành một lớp duy nhất để tiện thao tác rồi. Chúng ta sẽ làm:

1. Đặt kích thước của mảng (chúng ta sẽ sử dụng 16 vì nó được sử dụng rộng rãi trong bài báo và cho ViT-Base).
2. Lấy một hình ảnh đơn, in kích thước của nó và lưu trữ chiều cao và chiều rộng của nó.
3. Thêm một chiều batch vào hình ảnh đơn để nó tương thích với lớp PatchEmbedding của chúng ta.
4. Tạo một lớp (lớp chúng ta đã tạo ở phần 5.1.6) với patch_size=16 và embedding_dim=768 (từ Bảng 1 cho ViT-Base).
5. Đưa hình ảnh đơn qua lớp ở bước 4 để tạo một chuỗi các mảng nhúng.
6. Tạo một mảng nhúng lớp như trong phần 5.1.7.
7. Chèn thêm mảng nhúng lớp vào đầu chuỗi các mảng nhúng được tạo ở bước 5.
8. Tạo một mảng nhúng vị trí như trong phần 5.1.8.
9. Cộng mảng nhúng vị trí với mảng nhúng lớp và các mảng nhúng được tạo ở bước 7.

Code thôi nào:

set_seed(42)
#1. Set patch size
patch_size=16
#2. Print shape of original image tensor and get the image dimensions
print(f"Image tensor shape: {image.shape}")
height,width=image.shape[1],image.shape[2]
#3. Get image tensor and add batch dimension
x=image.unsqueeze(0)
print(f"Input image with batch dimension shape: {x.shape}")
#4. Create patch embedding layer
patch_embedding_layer=PatchEmbedding(in_channel=3,patch_size=patch_size,embedding_dim=768)
#5. Pass image through patch embedding layer
patch_embedding=patch_embedding_layer(x)
print(f"Patching embedding shape: {patch_embedding.shape}")
#6. Create class token embedding
batch_size=patch_embedding.shape[0]
embedding_dimension=patch_embedding.shape[-1]
class_token=nn.Parameter(torch.ones(batch_size,1,embedding_dimension),requires_grad=True)
print(f"Class token embedding shape: {class_token.shape}")
#7. Prepend class token embedding to patch embedding
patch_embedding_class_token=torch.cat((class_token,patch_embedding),dim=1)
print(f"Patch embedding with class token shape :{patch_embedding_class_token.shape}")

#8. Create postion embedding
number_of_patches=int((height*width)/(patch_size**2))
position_embedding=nn.Parameter(torch.ones(1,number_of_patches+1,embedding_dimension),requires_grad=True)
#9. Add postion embedding to patch embedding with class token
patch_and_position_embedding=patch_embedding_class_token+position_embedding
print(f"Patch and position embedding shape: {patch_and_position_embedding.shape}")

Output:

Image tensor shape: torch.Size([3, 224, 224])
Input image with batch dimension shape: torch.Size([1, 3, 224, 224])
Patching embedding shape: torch.Size([1, 196, 768])
Class token embedding shape: torch.Size([1, 1, 768])
Patch embedding with class token shape :torch.Size([1, 197, 768])
Patch and position embedding shape: torch.Size([1, 197, 768])

Những phần này tương ứng với đoạn code, bạn có thể nhìn hình phía dưới:

Source: Mrdbourke

Toàn bộ những phần mà chúng ta đã làm với công thức 1, cách hoạt động ra sao bạn có thể xem tại GIF phía dưới:

Source: Mrdbourke

5.2.1.Công thức 2: Multi-Head Attention (MSA)

Nhớ lại công thức 2. Tác giả đã viết:

$$\begin{aligned} \mathbf{z}_{\ell}^{\prime} &=\operatorname{MSA}\left(\operatorname{LN}\left(\mathbf{z}_{\ell-1}\right)\right)+\mathbf{z}_{\ell-1}, & & \ell=1 \ldots L \end{aligned}$$

Bạn có thấy một lớp MSA(Multi Head Self Attenion) nằm ngoài, và chứa một lớp LN(LayerNorm). Chúng ta sẽ gọi khối này là MSA.

Trái: Hình 1 từ bài báo ViT với các lớp Multi-Head Attention và Norm cũng như kết nối tắt (+) được làm nổi bật trong khối Encoder Transformer. Phải: Ánh xạ lớp Multi-Head Self Attention (MSA), lớp Norm và tắt đến các phần tương ứng của công thức 2 trong bài báo ViT. Source: Mrdbourke

5.2.2. Lớp LayerNorm(LN)

LayerNorm( hoặc Norm hoặc LN) là chuẩn hoá đầu vào trên chiều cuối cùng torch.nn.LayerNorm(). Bạn có thể tìm hiểu thêm vef Norm tại đây Sự khác biệt giữa các Norm có thể nhìn hình bên dưới:

Normalization. Source: Medium

5.2.3. Lớp Multi-Head Self Attention (MSA)

Để tìm hiểu thêm về Attention bạn có thể đọc paper Attention is all you need. Ban đầu Self Attention được thiết kế cho NLP với các task văn bản, Self Attention ấy một chuỗi các từ và sau đó tính toán từ nào nên chú ý nhiều hơn đến một từ khác. Nhưng đầu vào của chúng ta là một chuỗi các Patch Embedding, Self-attention và Multi-head attention sẽ tính toán các Patch Embedding nào của một bức ảnh có liên quan nhất tới các Patch Embedding khác, cuối cùng tạo thành một biểu diễn học được của bức ảnh. Bạn có thể tìm thấy định nghĩa về MSA của ViT Paper được định nghĩa trong phụ lục A.

Trái:Tổng quan về kiến trúc Vision Transformer từ Hình 1 của bài báo ViT. Bên phải: Định nghĩa của công thức 2, phần 3.1 và Phụ lục A của bài báo ViT được tô sáng để phản ánh các phần tương ứng của chúng trong Hình 1. Source Mrdbourke

Có tham số chúng ta quan tâm là Q,K,V, lần lượt là truy vấn( Queries), khoá (Key) và Giá trị(Values). là nền tảng Self Attention. Trong mô hình của chúng ta, chúng ta sẽ lấy 3 đầu ra của lớp LayerNorm phía trên và chuyển vào lớp này. Chúng ta có thể triển khai các Layer MSA trong Pytorch bằng các dùng torch.nn.MultiheadAttention(), các tham số cần quan tâm :

• embed_dim - Embedding Dimension từ Bảng 1 (Kích thước ẩn D).
• num_heads - Số lượng đầu Attention cần dùng , giá trị này cũng nằm trong Bảng 1 (Heads).
• Dropout - Bỏ ngẫu nhiên một số nút mạng trong quá trình huấn luyện.
• batch_first - Cho phép chỉ định dạng của tensor đầu vào và đầu ra. Nếu batch_first là True, thì tensor đầu vào và đầu ra sẽ có dạng (batch, seq, feature). Nếu batch_first là False (mặc định), thì tensor đầu vào và đầu ra sẽ có dạng (seq, batch, feature).

5.2.4. Gộp tất cả thành 1 lớp MSA

Vậy là lí thuyết đã rõ, chúng ta cần phải làm một số công việc:

1. Tạo một lớp có tên là MultiheadSelfAttentionBlock, kế thừa từ lớp torch.nn.Module.
2. Khởi tạo lớp với các tham số siêu (hyperparameters) từ Bảng 1 của bài báo ViT cho mô hình ViT-Base.
3. Tạo một lớp chuẩn hóa lớp (LN) với tham số normalized_shape bằng với kích thước nhúng (embedding dimension) D từ Bảng 1 (sử dụng torch.nn.LayerNorm()).
4. Tạo một lớp chú ý đa đầu (MSA) với các tham số phù hợp là embed_dim, num_heads, dropout và batch_first.
5. Tạo một phương thức cho lớp của chúng ta, truyền đầu vào qua lớp LN và lớp MSA (sử dụng phương thức forward()).

#1. Tạo lớp kế thừa nn.Module
class MultiheadSelfAttentionBlock(nn.Module):
    #2. Chỉnh tham số tương tự bảng 1
    def __init__(self,embedding_dim:int=768,# Hidden size D from Table 1 for ViT-Base
                 num_head:int=12,# Heads from Table 1 for ViT-Base
                 attn_dropout:float=0):
        super().__init__()
        #3.Tạo lớp LN
        self.layer_norm=nn.LayerNorm(normalized_shape=embedding_dim)
        #4. Tạo lớp MSA
        self.multihead_attn=nn.MultiheadAttention(embed_dim=embedding_dim,num_heads=num_head,dropout=attn_dropout,batch_first=True) 
    def forward(self,x):
        x=self.layer_norm(x)
        attn_output, _ = self.multihead_attn(query=x, # query embeddings
                                             key=x, # key embeddings
                                             value=x, # value embeddings
                                             need_weights=False)
        return attn_output

Thử xem lớp có hoạt động tốt hay không: Vậy là nó đã hoạt động.

5.3.1. Lớp MLP ( Multilayer Perceptron)

Phương trình 3:

$$\begin{aligned} \mathbf{z}_{\ell} &=\operatorname{MLP}\left(\operatorname{LN}\left(\mathbf{z}_{\ell}^{\prime}\right)\right)+\mathbf{z}_{\ell}^{\prime}, & & \ell=1 \ldots L \end{aligned}$$

Vậy là MLP sẽ bọc trong lớp LN và bổ sung ở cuối là kết nối tắt. Bạn có thể xem ảnh sau:

Trái: Hình 1 từ bài báo ViT với các lớp MLP và Norm cũng như kết nối dư (+) được làm nổi bật trong khối mã hóa Transformer. Phải: Ánh xạ lớp đa lớp perceptron (MLP), lớp Norm (LN) và kết nối dư đến các phần tương ứng của phương trình 3 trong bài báo ViT. Source Mrdbourke

5.3.2. Các lớp MLP

Multilayer perceptron (MLP) là một mô hình mạng nơ-ron nhiều tầng, có thể học được các hàm phi tuyến tính đối với dữ liệu phức tạp. MLP có thể được sử dụng cho các bài toán hồi quy hoặc phân loại. Trong ViT paper, tác giả có viết

The MLP contains two layers with a GELU non-linearity.

Khối MLP gồm hai tầng tuyến tính (torch.nn.Linear() trong PyTorch) và một hàm kích hoạt phi tuyến GELU (torch.nn.GELU() trong PyTorch). Tác giả cũng nói rằng mỗi tầng tuyến tính có một tầng dropout (torch.nn.Dropout() trong PyTorch) để giảm thiểu hiện tượng quá khớp. Các giá trị của các tham số cho các tầng tuyến tính có thể tìm thấy trong Bảng 1 của bài báo. Cuối cùng, tác giả cho biết cấu trúc của khối MLP như sau:

layer norm -> linear layer -> non-linear layer -> dropout -> linear layer -> dropout

5.3.3. Gộp thành lớp MLP

Công việc của chúng ta:

1. Tạo một lớp có tên là MLPBlock, kế thừa từ lớp torch.nn.Module.
2. Khởi tạo lớp với các tham số siêu (hyperparameters) từ Bảng 1 và Bảng 3 của bài báo ViT cho mô hình ViT-Base.
3. Tạo một lớp chuẩn hóa tầng (LN) với tham số normalized_shape bằng với kích thước nhúng (embedding dimension) D từ Bảng 1 (sử dụng torch.nn.LayerNorm()).
4. Tạo một chuỗi tuần tự các lớp MLP (s) sử dụng torch.nn.Linear(), torch.nn.Dropout() và torch.nn.GELU() với các giá trị tham số phù hợp từ Bảng 1 và Bảng 3.
5. Tạo một phương thức cho lớp của chúng ta, truyền đầu vào qua lớp LN và lớp MLP (s) (sử dụng phương thức forward()).

#1. Tạo lớp
class MLPBlock(nn.Module):
    def __init__(self,embedding_dim:int=768 # Hidden Size D from Table 1 for ViT-Base
                 ,mlp_size:int=3072 # MLP size from Table 1 for ViT-Base
                 ,drop_out:float=0.1): # Dropout from Table 3 for ViT-Base)
        super().__init__()
        #3. Tạo lớp LN
        self.layer_norm=nn.LayerNorm(normalized_shape=embedding_dim)
        #4. Tạo lớp MLP
        self.mlp=nn.Sequential(nn.Linear(in_features=embedding_dim,out_features=mlp_size),
                               nn.GELU(),
                               nn.Dropout(p=drop_out),
                               nn.Linear(in_features=mlp_size,out_features=embedding_dim),
                               nn.Dropout(p=drop_out))
        
    def forward(self,x):
        x=self.layer_norm(x)
        x=self.mlp(x)
        return x

Vậy là đã tạo xong, kiểm tra mô hình

mlp_block=MLPBlock(embedding_dim=768,mlp_size=3072,drop_out=0.1)
patched_image_through_mlp_block=mlp_block(patched_image_through_msa_block)
print(f"Input shape of MLP block: {patched_image_through_msa_block.shape}")
print(f"Output shape of MLP block: {patched_image_through_mlp_block.shape}")

Output: Input shape of MLP block: torch.Size([1, 197, 768])
        Output shape of MLP block: torch.Size([1, 197, 768])

Vậy là code đã hoạt động

5.4.1 Tạo Transformer Encoder

Vậy là chúng ta đã đi được 50% quãng đường. Tiếp tới chúng tôi sẽ cần xếp chồng khối MSA, MLP và tạo Transformer Encoder của mô hình ViT. Đoạn 4 phần 3.1 paper ViT có viết:

The Transformer encoder (Vaswani et al., 2017) consists of alternating layers of multiheaded selfattention (MSA, see Appendix A) and MLP blocks (Eq. 2, 3). Layernorm (LN) is applied before every block, and residual connections after every block (Wang et al., 2019; Baevski & Auli, 2019).

Vậy là chúng ta cần phải tạo các kết nôi còn lại. Các kết nối còn lại được lại các kết nối bỏ qua, lần đầu tiên giới thiệu trong bài báo Deep Residual Learning for Image Recognition. Nó cũng là trọng tâm của mạng phần dư Resnet. Trong Pseudo code có thể biểu diễn như sau:

x_input -> MSA_block -> [MSA_block_output + x_input] -> MLP_block -> [MLP_block_output + MSA_block_output + x_input] -> ...

5.4.2. Tạo Transformer Encoder bằng tay (Sử dụng các lớp có sẵn phía trên)

Chúng ta sẽ phải làm những việc sau:

1. Tạo một lớp có tên TransformerEncoderBlock kế thừa từ torch.nn.Module.
2. Khởi tạo lớp với các tham số siêu từ Bảng 1 và Bảng 3 của bài báo ViT cho mô hình ViT-Base.
3. Khởi tạo một khối MSA cho công thức 2 sử dụng MultiheadSelfAttentionBlock của chúng tôi từ mục 5.2.4 với các tham số thích hợp.
4. Khởi tạo một khối MLP cho công thức 3 sử dụng MLPBlock của chúng tôi từ mục 5.3.3 với các tham số thích hợp.
5. Tạo một phương thức forward() cho lớp TransformerEncoderBlock của chúng tôi.
6. Tạo một kết nối dư cho khối MSA (cho công thức 2).
7. Tạo một kết nối dư cho khối MLP (cho công thức 3).

Code:

#1. Tạo một lớp có tên TransformerEncoderBlock kế thừa từ torch.nn.Module. 
class TransformerEncoderBlock(nn.Module):
    #2. Khởi tạo lớp với các tham số siêu từ Bảng 1 và Bảng 3 của bài báo ViT cho mô hình ViT-Base. 
    def __init__(self,embedding_dim:int=768,num_heads:int=12,mlp_size:int=3072,mlp_dropout:float=0.1,attn_dropout:float=0):
        super().__init__()
        #3. Khởi tạo một khối MSA
        self.msa_block=MultiheadSelfAttentionBlock(embedding_dim=embedding_dim,num_head=num_heads,attn_dropout=attn_dropout)
        #4.Khởi tạo một khối MLP
        self.mlp_block=MLPBlock(embedding_dim=embedding_dim,mlp_size=mlp_size,drop_out=mlp_dropout)
    
    #5. Tạo 1 phương thức forward
    def forward(self,x):
        #6. Tạo kết nối bỏ qua cho khối MSA Block(Thêm Input vào output)
        x=self.msa_block(x)+x
        #7. Tạo kết nối bỏ qua cho khối MLP(Thêm Input vào output)
        x=self.mlp_block(x)+x
        return x

Hãy kiểm tra thử xem nó có hoạt động hay không:

mlp_block = MLPBlock(embedding_dim=768, # từ Bảng 1
                     mlp_size=3072, # từ Bảng 1
                     dropout=0.1) # từ Bảng 3

patched_image_through_mlp_block = mlp_block(patched_image_through_msa_block)
print(f"Input shape of MLP block: {patched_image_through_msa_block.shape}")
print(f"Output shape MLP block: {patched_image_through_mlp_block.shape}")

Chúng ta xem Transformer Encoder sau khi gộp MSA và MLP

5.4.3. Tạo Transformer Encoder với Pytorch's Transformer Layers

Chúng ta cũng có thể dùngtorch.nn.TransformerEncoderLayer để tạo Transformer Encoder. Code

torch_transformer_encoder_layer=nn.TransformerEncoderLayer(d_model=768 # Kích thước ẩn D từ bảng 1 ViT-Base
                                                           ,nhead=12 # Số đầu từ bảng 1 ViT-Base
                                                           ,dim_feedforward=3072 # Kích thước MLP
                                                           ,dropout=0.1# Lượng DropOut của các lớp dày đặc từ Bảng 3 mô hình ViT-Base
                                                           ,activation="gelu" # Hàm kích hoạt phi tuyến tính GeLU
                                                           ,batch_first=True #Dùng Batch_first 
                                                           ,norm_first=True)# Chuẩn hoá sau mỗi MSA và MLP

Thử hiển thị kiến trúc của mô hình

Giống với mô hình chúng ta đã tạo ở trên, đúng chứ?

5.5.1. Kết hợp tất cả các lớp lại với nhau tạo ViT Model

Vậy là chúng ta đã trải qua 3 phương trình là 1,2,3 và còn một phương trình cuối cùng là phương trình 4. Phương trình 4 có dạng như sau:

$$\begin{aligned} \mathbf{y} &=\operatorname{LN}\left(\mathbf{z}_{L}^{0}\right) & & \end{aligned}$$

Để viết phương trình này có thể sử dụng torch.nn.LayerNorm() và torch.nn.Linear() . Hãy bắt đầu viết thôi. Công việc mà ta sẽ phải làm:

1. Tạo một lớp có tên là ViT kế thừa từ torch.nn.Module.
2. Khởi tạo lớp với các tham số siêu từ Bảng 1 và Bảng 3 của bài báo ViT cho mô hình ViT-Base.
3. Đảm bảo kích thước ảnh chia hết cho Patch Size (ảnh nên được chia thành các đốm đều nhau).
4. Tính số lượng patch bằng công thức $N=H W / P^{2}$, trong đó H là chiều cao ảnh, W là chiều rộng ảnh và P là Patch Size.
5. Tạo một Class Token Embedding có thể học được (phương trình 1) như đã làm ở phần 5.1.7 ở trên.
6. Tạo một vector nhúng vị trí có thể học được (phương trình 1) như đã làm ở phần 5.1.8 ở trên.
7. Thiết lập lớp dropout Embedding.
8. Tạo lớp nhúng đốm bằng cách sử dụng lớp PatchEmbedding như đã làm ở phần 5.1.6 ở trên.
9. Tạo một chuỗi các khối mã hóa Transformer bằng cách truyền một danh sách các TransformerEncoderBlock được tạo ở phần 5.4.2 vào torch.nn.Sequential() (phương trình 2 và 3).
10. Tạo đầu MLP (còn gọi là đầu phân loại hoặc phương trình 4) bằng cách truyền một lớp torch.nn.LayerNorm() và một lớp torch.nn.Linear(out_features=num_classes) (trong đó num_classes là số lượng lớp mục tiêu) vào torch.nn.Sequential().
11. Tạo một phương thức nhận đầu vào là forward().
12. Lấy kích thước batch của đầu vào (chiều đầu tiên của hình dạng).
13. Tạo Patch Embedding bằng cách sử dụng lớp được tạo ở bước 8 (phương trình 1).
14. Tạo class token embedding bằng cách sử dụng lớp được tạo ở bước 5 và mở rộng nó trên số lượng batch được tìm thấy ở bước 11 bằng cách sử dụng torch.Tensor.expand() (phương trình 1).
15. Nối class token embedding được tạo ở bước 13 vào chiều đầu tiên của nhúng đốm được tạo ở bước 12 bằng cách sử dụng torch.cat() (phương trình 1).
16. Cộng position embedding được tạo ở bước 6 với Patch và Class token embedding được tạo ở bước 14 (phương trình 1).
17. Truyền patch và position embedding qua lớp dropout được tạo ở bước 7.
18. Truyền patch và position embedding từ bước 16 qua ngăn xếp các lớp mã hóa Transformer được tạo ở bước 9 (phương trình 2 và 3).
19. Truyền chỉ số 0 của đầu ra của ngăn xếp các lớp mã hóa Transformer từ bước 17 qua đầu phân loại được tạo ở bước 10 (phương trình 4).

Let's code:

#1. Create a ViT class that inherits from nn.Module
class ViT(nn.Module):
    def __init__(self,img_size:int=224,in_channels:int=3,patch_size:int=16,num_transformer_layers:int=12,embedding_dim:int=768,mlp_size:int=3072,num_heads:int=12,attn_dropout:float=0,mlp_dropout:float=0.1, embedding_dropout:float=0.1,num_classes:int=1000):
        super().__init__()
        #3. Make the image size is divisible by patch size
        assert img_size%patch_size==0, f"Image size must be divisible by patch size. Image size: {img_size}, Patch size: {patch_size}"
        # 4. Calculate the number of patches (height*width/patch^2)
        self.num_patches=(img_size*img_size)//(patch_size*patch_size)
        #5. Create learnable class embeddings 
        self.class_embedding=nn.Parameter(data=torch.randn(1,1,embedding_dim),requires_grad=True)
        # 6.Create learnable position embedding
        self.position_embedding=nn.Parameter(data=torch.randn(1, self.num_patches+1, embedding_dim),
                                               requires_grad=True)
        #7.Create embedding dropput values
        self.embedding_dropout=nn.Dropout(p=embedding_dropout)
        #8. Create patch embedding layer
        self.patch_embedding=PatchEmbedding(in_channel=in_channels,patch_size=patch_size,embedding_dim=embedding_dim)
        #9. Create Transformer Encoder blocks
        self.transformer_encoder=nn.Sequential(*[TransformerEncoderBlock(embedding_dim=embedding_dim,
                                                                            num_heads=num_heads,
                                                                            mlp_size=mlp_size,
                                                                            mlp_dropout=mlp_dropout) for _ in range(num_transformer_layers)])
        # 10. Create classifier head
        self.classifier = nn.Sequential(
            nn.LayerNorm(normalized_shape=embedding_dim),
            nn.Linear(in_features=embedding_dim,
                      out_features=num_classes))
    def forward(self,x):
        # 12. Get batch size
        batch_size=x.shape[0]
        #13. Create class token embedding and expand it to match the batchsize
        class_token=self.class_embedding.expand(batch_size,-1,-1)
        #14. Create patch embedding (Eq1)
        x=self.patch_embedding(x)
        #15. Concatenate class embedding and patch embedding (Eq1)
        x=torch.cat((class_token,x),dim=1)
        #16. Add position embedding to patch embedding(Eq1)
        x=self.position_embedding+x
        #17. Run embedding dropout (Appendix B1)
        x=self.embedding_dropout(x)
        #18. Pass patch,position and class embedding through transformer encoding
        x=self.transformer_encoder(x)
        #19. Put 0 index logts through classfier (Eq4)
        x=self.classifier(x[:,0])
        return x

Vậy là chúng ta đã xây dựng xong kiến trúc ViT. Hãy tạo một bản Demo đê xem mô hình chúng ta làm gì nào:

Vậy là chúng đã hoạt động .

5.6. Xem thông tin mô hình ViT mà chúng ta đã tạo

Chúng ta có thể dùng torchinfo

Khủng khiếp đúng không. 85,867,866 Params. Ngoài ViT Base , bạn có thể thay đổi để phù hợp như ViT Large hay ViT Huge. Lòng vòng quá rồi. Hãy đi đến phần cuối - Train Model!.

6. Train Model

Chúng ta sẽ cần thiết lập một số phần trước khi huấn luyện mô hình. Trong paper ViT, nhóm tác giả có viết

Chúng ta có thể sử dụng các tham số trong đoạn trên.

6.1. Tạo Optimizers

Nhìn ảnh trên, chúng ta có thể thấy họ đã chọn sử dụng trình tối ưu hóa "Adam" thay vì SGD. Nhóm tác giả cũng đặt ra tham số $\beta$ (beta) của Adam với $\beta_{1}=0.9, \beta_{2}=0.999$.

Ngoài ra nhóm tác giả cũng cài đặt weight_decay trong quá trình tối ưu hoá để tránh overfitting, chúng ta có thể cài đặt chúng dựa vào mô hình đã huấn luyện trên ImageNet 1k với weight_decay=0.3 Chúng ta có để code đoạn này:

optimizer = torch.optim.Adam(params=vit.parameters(),
                             lr=3e-3, # Base LR từ bảng 3 cho ImageNet-1k
                             betas=(0.9, 0.999), 
                             weight_decay=0.3)
                         

6.2. Tạo hàm mất mát (Loss Function)

Phần này chúng ta sẽ sử dụng hàm CrossEntropyLoss để tính toán Loss Function. Code:

loss_fn = torch.nn.CrossEntropyLoss()

6.3. Thiết lập EarlyStopping

Mục đích để tracking lại hiệu suất của mô hình. Rồi có quyết dịnh dừng sớm để tránh lãng phí tài nguyên hay không. Code như sau:

import numpy as np
class EarlyStopping(object):
    def __init__(self, mode='min', min_delta=0, patience=10, percentage=False):
        self.mode = mode
        self.min_delta = min_delta
        self.patience = patience
        self.best = None
        self.num_bad_epochs = 0
        self.is_better = None
        self._init_is_better(mode, min_delta, percentage)

        if patience == 0:
            self.is_better = lambda a, b: True
            self.step = lambda a: False

    def step(self, metrics):
        if self.best is None:
            self.best = metrics
            return False

        if np.isnan(metrics):
            return True

        if self.is_better(metrics, self.best):
            self.num_bad_epochs = 0
            self.best = metrics
            print('improvement!')
        else:
            self.num_bad_epochs += 1
            print(f'no improvement, bad_epochs counter: {self.num_bad_epochs}')

        if self.num_bad_epochs >= self.patience:
            return True

        return False

    def _init_is_better(self, mode, min_delta, percentage):
        if mode not in {'min', 'max'}:
            raise ValueError('mode ' + mode + ' is unknown!')
        if not percentage:
            if mode == 'min':
                self.is_better = lambda a, best: a < best - min_delta
            if mode == 'max':
                self.is_better = lambda a, best: a > best + min_delta
        else:
            if mode == 'min':
                self.is_better = lambda a, best: a < best - (
                            best * min_delta / 100)
            if mode == 'max':
                self.is_better = lambda a, best: a > best + (
                            best * min_delta / 100)

6.4. Thiết lập hàm để train model

Chúng ta có thể thiết lập hàm train model với 3 thành phần chính sau:

• train_step: Thực hiện bước huấn luyện mô hình trên một batch dữ liệu từ train dataloader. Hàm này nhận vào các tham số là mô hình, dataloader, hàm mất mát và bộ tối ưu hóa. Hàm này trả về giá trị độ chính xác và mất mát trên batch đó.
• test_step: Tương tự nhưng trên testdataloaders
• train: Kích hoạt 2 hàm phía trên

Bạn cũng thể tinh chỉnh tên của dự án trên Wandb bằng cách thay đổi run=wandb.init(project="Vision Transformer Plane Classification Model") thành tên mà bạn muốn.

Chúng ta có thể code như sau:

import torch
import torch.nn as nn
from tqdm.auto import tqdm
from typing import List,Tuple,Dict

def train_step(model:torch.nn.Module,dataloader:torch.utils.data.DataLoader,loss_fn:torch.nn.Module,optimizers:torch.optim.Optimizer):
    wandb.watch(model, log_freq=100)
    model.train()
    train_acc,train_loss=0,0
    for batch,(X,y) in enumerate(dataloader):
        X,y=X.to(devices),y.to(devices)
        y_pred=model(X)
        loss=loss_fn(y_pred,y)
        train_loss+=loss.item()
        optimizers.zero_grad()
        loss.backward()
        optimizers.step()
        y_pred_class=torch.argmax(torch.softmax(y_pred,dim=1),dim=1)
        train_acc +=(y_pred_class==y).sum().item()/len(y_pred)
    train_acc/=len(dataloader)
    train_loss/=len(dataloader)
    return train_acc,train_loss

def test_step(model:torch.nn.Module,dataloader:torch.utils.data.DataLoader,loss_fn:torch.nn.Module):
    model.eval()
    test_loss_values,test_acc_values=0,0
    with torch.inference_mode():
        for batch,(X,y) in enumerate(dataloader):
            X,y=X.to(devices),y.to(devices)
            y_test_pred_logits=model(X)
            
            test_loss=loss_fn(y_test_pred_logits,y)
            test_loss_values+=test_loss.item()
            
            y_pred_class=torch.argmax(y_test_pred_logits,dim=1)
            test_acc_values += ((y_pred_class==y).sum().item()/len(y_test_pred_logits))
        test_loss_values/=len(dataloader)
        test_acc_values/=len(dataloader)
    return test_loss_values,test_acc_values
run=wandb.init(project="Vision Transformer Plane Classification Model")
def train(model: torch.nn.Module, train_dataloader: torch.utils.data.DataLoader, test_dataloader: torch.utils.data.DataLoader, optimizer: torch.optim.Optimizer, loss_fn: torch.nn.Module = nn.CrossEntropyLoss(), epochs: int = 100, early_stopping=None):
    result = {
        "train_loss": [],
        "train_acc": [],
        "test_loss": [],
        "test_acc": []
    }

    for epoch in tqdm(range(epochs)):
        train_acc, train_loss = train_step(model=model, dataloader=train_dataloader, loss_fn=loss_fn, optimizers=optimizer)
        test_loss, test_acc = test_step(model=model, dataloader=test_dataloader, loss_fn=loss_fn)
        
        print(
            f"Epoch: {epoch + 1} | "
            f"train_loss: {train_loss:.4f} | "
            f"train_acc: {train_acc:.4f} | "
            f"test_loss: {test_loss:.4f} | "
            f"test_acc: {test_acc:.4f}"
        )

        # Update results dictionary
        result["train_loss"].append(train_loss)
        result["train_acc"].append(train_acc)
        result["test_loss"].append(test_loss)
        result["test_acc"].append(test_acc)
        wandb.log({"Train Loss": train_loss,
                   "Test Loss": test_loss,
                   "Train Accuracy": train_acc,
                   "Test Accuracy": test_acc,"Epoch":epoch})
        # Check for early stopping
        if early_stopping is not None:
            if early_stopping.step(test_loss):  # You can use any monitored metric here
                print(f"Early stopping triggered at epoch {epoch + 1}")
                break

    return result

6.4. Fit Model

Giờ chúng ta có thể Fit những dữ liệu chúng ta có vào mô hình. Chúng ta cần truyền vào model, train_dataloader, test_dataloader, optimizer và loss_function, early_stopping nếu không dùng có thể để None.

early_stopping = EarlyStopping(mode='min', patience=10)
devices="cuda" if torch.cuda.is_available() else "cpu"
model_result=train(model=vit,train_dataloader=train_dataloaders,test_dataloader=test_dataloader,optimizer=optimizer,loss_fn=loss_fn,epochs=60,early_stopping=early_stopping)
run.finish() #finish wandb

6.5. Tạo hàm save model

Sau khi train model, chúng ta có thể lưu mô hình vào máy để lần sau có thể dùng tiếp.

import torch
from pathlib import Path

def save_model(model:torch.nn.Module,target_dir:str,model_name:str):
    target_dir_path = Path(target_dir)
    target_dir_path.mkdir(parents=True,
                        exist_ok=True)

  # Create model save path
    assert model_name.endswith(".pth") or model_name.endswith(".pt"), "model_name should end with '.pt' or '.pth'"
    model_save_path = target_dir_path / model_name
    
      # Save the model state_dict()
    print(f"[INFO] Saving model to: {model_save_path}")
    torch.save(obj=model.state_dict(),
                 f=model_save_path)

Dùng như sau:

save_model(model=vit,
              target_dir="models",
              model_name="ViT_for_Classification.pt")

7. Kết quả

Mô hình chúng ta đã chạy. Nhưng tại sao kết quả lại tệ, phải chăng chúng ta đã bỏ lỡ bước gì đó. Có vài nguyên nhân khiến mô hình của chúng ta không được tốt. Hãy so sánh:

Giá trị siêu tham số	Bài báo ViT	Bản triển khai của chúng ta
Số lượng hình ảnh huấn luyện	1.3 triệu (ImageNet-1k), 14 triệu (ImageNet-21k), 303 triệu (JFT)	Hơn 5500
Số epochs	7 (cho bộ dữ liệu lớn nhất), 90, 300 (cho ImageNet)	10
Kích thước batch	4096	32
Tăng tốc độ học	10k bước (Bảng 3)	Không
Giảm tốc độ học	Tuyến tính / Cosine (Bảng 3)	Không
Gradient Clipping	Global norm 1 (Bảng 3)	Không

Bài pretrain cho ViT models đã được cập nhật, bạn có thể vào link sau: https://viblo.asia/p/pretrain-model-vision-transformer-in-pytorch-GAWVpyna405

References

1. Pytorch Tutorial: https://www.learnpytorch.io/08_pytorch_paper_replicating/#9-setting-up-training-code-for-our-vit-model
2. Paper ViT: https://arxiv.org/abs/2010.11929
3. Paper ResidualNet: https://arxiv.org/abs/1512.03385v1
4. Paper Transformer: https://arxiv.org/abs/1706.03762
5. Wandb: https://wandb.ai/home
6. Full Source code: https://www.kaggle.com/tnguynfew/vit-for-animal-classification

Cảm ơn đã đọc bài này của mình. Nếu bạn các bạn thấy hữu ích có thể cho mình xin 1 upvote.