OpenSora代码理解（一）Video Embedding

Video Embedding

class PatchEmbed3D(nn.Module):
    """Video to Patch Embedding.

    Args:
        patch_size (int): Patch token size. Default: (2,4,4).
        in_chans (int): Number of input video channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(
        self,
        patch_size=(2, 4, 4),
        in_chans=3,
        embed_dim=96,
        norm_layer=None,
        flatten=True,
    ):
        super().__init__()
        self.patch_size = patch_size
        self.flatten = flatten

        self.in_chans = in_chans
        self.embed_dim = embed_dim

        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
        if norm_layer is not None:
            self.norm = norm_layer(embed_dim)
        else:
            self.norm = None

    def forward(self, x):
        """Forward function."""
        # padding
        _, _, D, H, W = x.size()
        if W % self.patch_size[2] != 0:
            x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2]))
        if H % self.patch_size[1] != 0:
            x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1]))
        if D % self.patch_size[0] != 0:
            x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0]))

        x = self.proj(x)  # (B C T H W)
        if self.norm is not None:
            D, Wh, Ww = x.size(2), x.size(3), x.size(4)
            x = x.flatten(2).transpose(1, 2)
            x = self.norm(x)
            x = x.transpose(1, 2).view(-1, self.embed_dim, D, Wh, Ww)
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)  # BCTHW -> BNC
        return x

这段代码定义了一个用于将视频数据划分为块并嵌入到高维空间的 PatchEmbed3D 类。下面我们会一步步分析 forward 方法中数据的维度变化，帮助你更好地理解每个操作如何改变数据的形状：

输入维度

输入 x 是一个五维张量，其形状为 (B, C, D, H, W)：

• B 是批次大小（batch size）。
• C 是通道数，对于视频通常是3（RGB）。
• D 是深度，代表视频的帧数或时间序列的长度。
• H 是每帧视频的高度。
• W 是每帧视频的宽度。

填充

为了确保每个维度（D, H, W）都可以被 patch_size 中对应的数值整除，进行以下填充操作：

1. 对宽度 W 的填充：

   • 如果 W % self.patch_size[2] 不为0，则在宽度的末尾进行填充，填充的大小为 self.patch_size[2] - W % self.patch_size[2]。
   • 形状变为 (B, C, D, H, W + padding)。

2. 对高度 H 的填充：

   • 如果 H % self.patch_size[1] 不为0，则在高度的末尾进行填充，填充的大小为 self.patch_size[1] - H % self.patch_size[1]。
   • 形状变为 (B, C, D, H + padding, W + padding)（前一步可能已经对 W 进行了填充）。

3. 对深度 D 的填充：

   • 如果 D % self.patch_size[0] 不为0，则在深度的末尾进行填充，填充的大小为 self.patch_size[0] - D % self.patch_size[0]。
   • 形状变为 (B, C, D + padding, H + padding, W + padding)。

卷积操作

• 使用一个 3D 卷积层 self.proj（具有 kernel 和 stride 都为 patch_size），此操作将每个 patch_size 块映射到 embed_dim 维空间中。
• 输出张量的形状变为 (B, embed_dim, D', H', W')，其中 D', H', W' 是卷积操作后的深度、高度和宽度，取决于 patch_size 和填充。

归一化和重排

• 如果设置了归一化层 self.norm，先将数据展平到 (B, embed_dim, D'*H'*W')，转置为 (B, D'*H'*W', embed_dim)，应用归一化，然后恢复为原来的形状 (B, embed_dim, D', H', W')。

扁平化操作

• 如果 self.flatten 为 True，则将输出张量扁平化为 (B, D'*H'*W', embed_dim)，即每个视频块转换成一个嵌入向量，以适应后续的处理或模型结构。

通过上述步骤，PatchEmbed3D 将输入的视频数据转换为一组高维表示，这些表示可用于视频分类、检测或其他视频理解任务。这个转换处理能够将视频中的局部空间和时间信息有效编码。