8.16笔记，VIT

import torch
import torch.nn as nn
import torch.nn.functional as F

def image2emb_naive(image, patch_size, weight):
    # image shape: bs*channel*h*w
    patch = F.unfold(image, kernel_size=patch_size, stride=patch_size).transpose(-1,-2)
    patch_embedding = patch @ weight
    return patch_embedding

def image2emb_conv(image, kernel, stride):
    conv_output = F.conv2d(image, kernel, stride=stride) #bs*oc*oh*ow
    bs, oc, oh, ow = conv_output.shape
    patch_embedding = conv_output.reshape(bs, oc, oh*ow).transpose(-1,-2)
    return patch_embedding

# test code for image2emb
bs, ic, image_h, image_w = 1,3,8,8
patch_size = 4
model_dim = 8
max_num_token = 16
num_classes = 10
label = torch.randint(10,(bs,))

patch_depth = patch_size*patch_size*ic
image = torch.randn(bs, ic, image_h, image_w)
weight = torch.randn(patch_depth, model_dim) #model_dim是输出通道数目，patch_depth是卷积核的面积乘以输入通道数

patch_embedding_naive = image2emb_naive(image, patch_size, weight) #分块方法得到embedding
kernel = weight.transpose(0,1).reshape(-1, ic, patch_size, patch_size) #oc*ic*kh*kv

patch_embedding_conv = image2emb_conv(image, kernel, patch_size) #二维卷积的方法得到embedding

print(patch_embedding_conv)
print(patch_embedding_naive)

## step2
cls_token_embedding = torch.randn(bs, 1, model_dim,requires_grad=True)
token_embedding = torch.cat([cls_token_embedding, patch_embedding_conv], dim=1)

## step3
position_embedding_table = torch.randn(max_num_token, model_dim, requires_grad=True)
seq_len = token_embedding.shape[1]
position_embedding = torch.tile(position_embedding_table[:seq_len],[token_embedding.shape[0],1,1])
token_embedding += position_embedding

## step4
encoder_layer = nn.TransformerEncoderLayer(d_model=model_dim, nhead=8)
transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
encoder_output = transformer_encoder(token_embedding)

## step5
cls_token_output = encoder_output[:, 0, :]
linear_layer = nn.Linear(model_dim, num_classes)
logits = linear_layer(cls_token_output)
loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(logits, label)
print(loss)

tensor([[[ -3.6616,   1.2795,   8.1758,   3.0673,  -7.8565,  -5.1824, -13.5376,
           -0.1308],
         [  6.9238,  -7.3512,  -1.2001,  -4.0659,  -2.5937,  -5.9933, -10.0656,
            0.6723],
         [ 11.3471,   7.7629,  14.7286,  -6.5203,  -9.3314,  -0.9771,  -3.2062,
           -5.0264],
         [  6.5818,   1.7215,  -3.0544,   5.1548,  14.1692,   8.5659,   5.4677,
            1.7021]]])
tensor([[[ -3.6616,   1.2795,   8.1758,   3.0673,  -7.8565,  -5.1824, -13.5376,
           -0.1308],
         [  6.9238,  -7.3512,  -1.2001,  -4.0659,  -2.5937,  -5.9933, -10.0656,
            0.6723],
         [ 11.3471,   7.7629,  14.7286,  -6.5203,  -9.3314,  -0.9771,  -3.2062,
           -5.0264],
         [  6.5818,   1.7215,  -3.0544,   5.1548,  14.1692,   8.5659,   5.4677,
            1.7021]]])
tensor(2.8527, grad_fn=<NllLossBackward0>)