transformer--ViT
代码为了处理二维图像,我们将尺寸为 H×W×C的图像reshape为拉平的2维图块,尺寸为 (N×(P^2×C))。其中, (P,P)为图块的大小, N=HW/P^2 。 N 是图块的数量,会影响输入序列的长度。Transformer在所有图层上使用恒定的隐矢量D,因此我们将图块拉平,并使用可训练的线性投影映射到D的大小,将此投影的输出称为patch embedding。对应代码如下:直接暴力拉伸# Transformer. n, h, w, c = x.shape x = jnp.r




