Vision Transformer中的数据形状梳理

关键零件内部数据形状

PatchEmbedding

原始输入x形状为(b, 3, 224, 224)
每一个patch的尺寸为(3, 16, 16)，则可以分割为1414个patches，于是先分割为(b, 1414, 31616)即(b, 196, 768)，一个patch_embedding的长度1024，则经过一个mlp，embedding为为(b, 196, 1024)，添加可学习的cls_token形状为(1, 1, 1024)作为最后送入分类mlp的cls向量，此时形状为(b, 197, 1024)，在此基础上添加绝对位置编码，直接加上去，保持形状不变。
得到输出x形状为(b, 1414+1, 316*16) -> (b, 197, 1024)

transformer encorder层

输入形状为 (b, 197, 1024)
self-attention层
设置头数8与每个头的dim数64，则qkv向量长度为512，输入数据送入mlp中得到输出形状(b,197,512*3)，拆开后得到Q,K,V矩阵为(b, 8 , 197, 64)，QK取内积得到(b, 8 , 197, 197)，softmax后与V取值得到(b, 8 , 197, 64)，整理为(b, 197, 512)，经过mlp输出为 (b, 197, 1024)
输出形状为 (b, 197, 1024)
FFN层
两个全连接
(b, 197, 1024) -> (b, 197,xxxx) -> (b, 197, 1024)
输出形状为 (b, 197, 1024)
即整个transformerencorder块不改变形状

整体结构

输入(b, 3, 224, 224)
Embedding编码
输出(b, 197, 1024)
经过数个transformerencorder块，形状不发生改变
输出(b, 197, 1024)
如果使用cls_token则取出其中cls_token(b, 1, 1024)转化为(b, 1024)
如果不适用，上述输出为(b, 196, 1024)，1,2维转置后进行全局平均池化为(b, 1024)
输出为(b, 1024)
经过全连接层获得最终输出
输出为(b, classes_num)