视频生成模子Sora的全面解析：从AI绘画、ViT到ViViT、DiT、VDT、NaViT、Vid - ToB企服应用市场:ToB评测及商务社交产业平台

起首定义ViViT类，且定义相关变量

# 定义ViViT模型类
class ViViT(nn.Module):
def __init__(self, image_size, patch_size, num_classes, num_frames, dim=192, depth=4, heads=3, pool='cls', in_channels=3, dim_head=64, dropout=0.,
emb_dropout=0., scale_dim=4):
super().__init__() # 调用父类的构造函数
# 检查pool参数是否有效
assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
# 确保图像尺寸能被patch尺寸整除
assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
# 计算patch数量
num_patches = (image_size // patch_size) ** 2
# 计算每个patch的维度
patch_dim = in_channels * patch_size ** 2
# 将图像切分成patch并进行线性变换的模块
self.to_patch_embedding = nn.Sequential(
Rearrange('b t c (h p1) (w p2) -> b t (h w) (p1 p2 c)', p1=patch_size, p2=patch_size),
nn.Linear(patch_dim, dim),
)

复制代码

为方便各人明白，我得表明一下上面中这行的寄义：

Rearrange('b t c (h p1) (w p2) -> b t (h w) (p1 p2 c)', p1=patch_size, p2=patch_size)

复制代码

且为方便各人和我之前介绍ViT的文章前后连贯起来，故还是用的ViT那篇文章中的例子(此文的第4部门)
以ViT_base_patch16为例，一张224 x 224的图片先分割成 16 x 16 的 patch ，很显然会因此而存在

个 patch
且图片的长宽由原来的224 x 224 变成：14 x 14(因为224/16 = 14)

16*16

...

以是对于上面那行意味着可以让批次大小b=1、时间维度t=2、RGB图像的通道数c=3
原始维度即为：
(1, 2, 3, 旧的长 = 224 patch_size = 16, 旧的宽 = 224 patch_size = 16)，Rearrange之后的维度则变为：
(1, 2, 新的长14 x 新的宽14 = 196, 16 x 16 x 3 = 768)

初始化位置编码和cls token
self.pos_embedding 的维度为(1, num_frames, num_patches + 1, dim)
在这里，num_frames 是 t，num_patches 是 n=196，dim 是 768，因此 pos_embedding 维度为 (1,2,197,768)

# 位置编码
self.pos_embedding = nn.Parameter(torch.randn(1, num_frames, num_patches + 1, dim))
# 空间维度的cls token
self.space_token = nn.Parameter(torch.randn(1, 1, dim))
# 空间变换器
self.space_transformer = Transformer(dim, depth, heads, dim_head, dim * scale_dim, dropout)
# 时间维度的cls token
self.temporal_token = nn.Parameter(torch.randn(1, 1, dim))
# 时间变换器
self.temporal_transformer = Transformer(dim, depth, heads, dim_head, dim * scale_dim, dropout)
# dropout层
self.dropout = nn.Dropout(emb_dropout)
# 池化方式
self.pool = pool
# 最后的全连接层，用于分类
self.mlp_head = nn.Sequential(
nn.LayerNorm(dim),
nn.Linear(dim, num_classes)
)

复制代码

patch嵌入和cls token的拼接
输入数据 x 的维度在颠末嵌入层后变为 (1,2,196,768)
self.space_token 的初始维度为 (1,1,768)，被复制扩展成 (1,2,1,768) 以匹配批次和时间维度
cls_space_tokens 和 x 在patch维度上拼接后，维度变为 (1,2,197,768)
为何拼接之后成197了呢？原因很简朴，如ViT那篇文章中所述：“[class] token的维度为 [1, 768] ，通过Concat操作，[196, 768] 与 [1, 768] 拼接得到 [197, 768]”

def forward(self, x):
# 将输入数据x转换为patch embeddings
x = self.to_patch_embedding(x)
b, t, n, _ = x.shape # 获取batch size, 时间维度, patch数量
# 在每个空间位置加上cls token
cls_space_tokens = repeat(self.space_token, '() n d -> b t n d', b=b, t=t)
x = torch.cat((cls_space_tokens, x), dim=2) # 在维度2上进行拼接

复制代码

添加位置编码和应用dropout
加上位置编码后，x 保持 (1,2,197,768) 维度不变。应用dropout后，x 的维度仍然不变

x += self.pos_embedding[:, :, :(n + 1)] # 加上位置编码
x = self.dropout(x) # 应用dropout

复制代码

空间Transformer
重排 x 的维度为 (2,197,768)，因为 b×t=1×2=2
空间Transformer处理后，x 的维度变为 (2,197,768)

# 将(b, t, n, d)重排为((b t), n, d)，为了应用空间变换器
x = rearrange(x, 'b t n d -> (b t) n d')
x = self.space_transformer(x) # 应用空间变换器
x = rearrange(x[:, 0], '(b t) ... -> b t ...', b=b) # 把输出重排回(b, t, ...)

复制代码

时间Transformer
self.temporal_token 的初始维度为(1,1,768)，被复制扩展成 (1,2,768)
cls_temporal_tokens 和 x 在时间维度上拼接后，维度变为(1,3,768)

# 在每个时间位置加上cls token
cls_temporal_tokens = repeat(self.temporal_token, '() n d -> b n d', b=b)
x = torch.cat((cls_temporal_tokens, x), dim=1) # 在维度1上进行拼接
x = self.temporal_transformer(x) # 应用时间变换器

复制代码

池化
如果 self.pool 是 'mean'，则对 x 在时间维度上取均值，结果维度变为 (1,768)
如果不是 'mean'，则直接取 x 的第一个时间维度的cls token，结果维度同样是 (1,768)

# 根据pool参数选择池化方式
x = x.mean(dim=1) if self.pool == 'mean' else x[:, 0]

复制代码

分类头
self.mlp_head，将 (1,768) 维度的 x 转换为终极的分类结果，其维度取决于类别数num_classes，如果 num_classes 是 10，则终极输出维度为 (1,10)

# 通过全连接层输出最终的分类结果
return self.mlp_head(x)

复制代码