CogVideo & CogVideoX 微调代码源码解析（七）
.\cogvideo-finetune\sat\sgm\modules\autoencoding\regularizers\base.py
# 导入抽象方法和类型注解
from abc import abstractmethod
# 导入任意类型和元组类型
from typing import Any, Tuple
# 导入 PyTorch 和功能模块
import torch
import torch.nn.functional as F
# 导入神经网络模块
from torch import nn
# 定义一个抽象正则化器类，继承自 nn.Module
class AbstractRegularizer(nn.Module):
# 初始化方法
def __init__(self):
# 调用父类的初始化方法
super().__init__()
# 定义前向传播方法，接受一个张量并返回张量和字典
def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, dict]:
# 抛出未实现错误，强制子类实现该方法
raise NotImplementedError()
# 定义获取可训练参数的抽象方法
@abstractmethod
def get_trainable_parameters(self) -> Any:
# 抛出未实现错误，强制子类实现该方法
raise NotImplementedError()
# 定义身份正则化器类，继承自 AbstractRegularizer
class IdentityRegularizer(AbstractRegularizer):
# 实现前向传播方法，接受一个张量并返回该张量和空字典
def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, dict]:
return z, dict()
# 实现获取可训练参数的方法，返回一个生成器
def get_trainable_parameters(self) -> Any:
# 生成器不返回任何值
yield from ()
# 定义测量困惑度的函数，接受预测索引和质心数量作为参数
def measure_perplexity(predicted_indices: torch.Tensor, num_centroids: int) -> Tuple[torch.Tensor, torch.Tensor]:
# src: https://github.com/karpathy/deep-vector-quantization/blob/main/model.py
# 评估集群的困惑度，当困惑度等于嵌入数量时，所有集群的使用是完全均匀的
# 将预测索引转化为独热编码并重塑为二维张量
encodings = F.one_hot(predicted_indices, num_centroids).float().reshape(-1, num_centroids)
# 计算每个质心的平均概率
avg_probs = encodings.mean(0)
# 计算困惑度
perplexity = (-(avg_probs * torch.log(avg_probs + 1e-10)).sum()).exp()
# 计算使用的集群数量
cluster_use = torch.sum(avg_probs > 0)
# 返回困惑度和集群使用数量
return perplexity, cluster_use
.\cogvideo-finetune\sat\sgm\modules\autoencoding\regularizers\finite_scalar_quantization.py
# Finite Scalar Quantization: VQ-VAE Made Simple - https://arxiv.org/abs/2309.15505
# 代码来自 Jax 版本的附录 A.1
from typing import List, Optional # 导入 List 和 Optional 类型以便进行类型注释
import torch # 导入 PyTorch 库
import torch.nn as nn # 导入 PyTorch 的神经网络模块
from torch.nn import Module # 从 nn 模块中导入 Module 基类
from torch import Tensor, int32 # 导入 Tensor 和 int32 类型
from torch.cuda.amp import autocast # 导入自动混合精度的上下文管理器
from einops import rearrange, pack, unpack # 从 einops 库导入重排、打包和解包函数
# helper functions
def exists(v): # 定义一个函数检查变量 v 是否存在
return v is not None # 返回 v 是否不为 None
def default(*args): # 定义一个函数以返回第一个存在的参数
for arg in args: # 遍历所有参数
if exists(arg): # 如果参数存在
return arg # 返回该参数
return None # 如果没有参数存在，返回 None
def pack_one(t, pattern): # 定义一个函数以打包一个张量 t
return pack([t], pattern) # 将 t 放入列表中并按照模式打包
def unpack_one(t, ps, pattern): # 定义一个函数以解包一个张量 t
return unpack(t, ps, pattern)[0] # 解包 t，返回第一个解包结果
# tensor helpers
def round_ste(z: Tensor) -> Tensor: # 定义一个函数以进行带有直通梯度的四舍五入
"""Round with straight through gradients.""" # 函数说明
zhat = z.round() # 对 z 进行四舍五入
return z + (zhat - z).detach() # 返回 z 加上 zhat 和 z 的差的梯度不跟随版本
# main class
class FSQ(Module): # 定义 FSQ 类，继承自 Module
def __init__( # 定义初始化方法
self,
levels: List[int], # 量化级别的列表
dim: Optional[int] = None, # 输入维度，可选
num_codebooks=1, # 码本数量，默认为 1
keep_num_codebooks_dim: Optional[bool] = None, # 是否保留码本维度的可选参数
scale: Optional[float] = None, # 缩放因子，可选
):
super().__init__() # 调用父类的初始化方法
_levels = torch.tensor(levels, dtype=int32) # 将 levels 转换为 int32 类型的张量
self.register_buffer("_levels", _levels, persistent=False) # 注册不持久化的缓冲区
_basis = torch.cumprod(torch.tensor([1] + levels[:-1]), dim=0, dtype=int32) # 计算基础张量
self.register_buffer("_basis", _basis, persistent=False) # 注册不持久化的缓冲区
self.scale = scale # 保存缩放因子
codebook_dim = len(levels) # 计算码本维度
self.codebook_dim = codebook_dim # 保存码本维度
effective_codebook_dim = codebook_dim * num_codebooks # 计算有效码本维度
self.num_codebooks = num_codebooks # 保存码本数量
self.effective_codebook_dim = effective_codebook_dim # 保存有效码本维度
keep_num_codebooks_dim = default(keep_num_codebooks_dim, num_codebooks > 1) # 确定是否保留码本维度
assert not (num_codebooks > 1 and not keep_num_codebooks_dim) # 确保规则有效
self.keep_num_codebooks_dim = keep_num_codebooks_dim # 保存是否保留的标志
self.dim = default(dim, len(_levels) * num_codebooks) # 设置输入维度
has_projections = self.dim != effective_codebook_dim # 检查是否需要投影
self.project_in = nn.Linear(self.dim, effective_codebook_dim) if has_projections else nn.Identity() # 输入投影层
self.project_out = nn.Linear(effective_codebook_dim, self.dim) if has_projections else nn.Identity() # 输出投影层
self.has_projections = has_projections # 保存是否有投影的标志
self.codebook_size = self._levels.prod().item() # 计算码本大小
implicit_codebook = self.indices_to_codes(torch.arange(self.codebook_size), project_out=False) # 获取隐式码本
self.register_buffer("implicit_codebook", implicit_codebook, persistent=False) # 注册隐式码本缓冲区
def bound(self, z: Tensor, eps: float = 1e-3) -> Tensor: # 定义边界函数
"""Bound `z`, an array of shape (..., d).""" # 函数说明
half_l = (self._levels - 1) * (1 + eps) / 2 # 计算半边界
offset = torch.where(self._levels % 2 == 0, 0.5, 0.0) # 计算偏移量
shift = (offset / half_l).atanh() # 计算偏移量的反双曲正切
return (z + shift).tanh() * half_l - offset # 返回边界调整后的结果
# 定义量化函数，将输入张量 z 进行量化，并返回量化后的 zhat，形状与 z 相同
def quantize(self, z: Tensor) -> Tensor:
"""Quantizes z, returns quantized zhat, same shape as z."""
# 对输入 z 进行边界处理后，执行四舍五入量化
quantized = round_ste(self.bound(z))
# 计算半宽度，用于重新归一化到 [-1, 1] 的范围
half_width = self._levels // 2 # Renormalize to [-1, 1].
# 返回量化后的值，归一化到 [-1, 1] 范围
return quantized / half_width
# 定义缩放和偏移函数，将归一化的 zhat 转换为相应的范围
def _scale_and_shift(self, zhat_normalized: Tensor) -> Tensor:
# 计算半宽度
half_width = self._levels // 2
# 将归一化的 zhat 进行缩放和偏移，返回相应值
return (zhat_normalized * half_width) + half_width
# 定义反缩放和偏移函数，将 zhat 转换回归一化值
def _scale_and_shift_inverse(self, zhat: Tensor) -> Tensor:
# 计算半宽度
half_width = self._levels // 2
# 进行反缩放和偏移，返回归一化值
return (zhat - half_width) / half_width
# 定义将代码转换为索引的函数
def codes_to_indices(self, zhat: Tensor) -> Tensor:
"""Converts a `code` to an index in the codebook."""
# 确保 zhat 的最后一维与代码本的维度匹配
assert zhat.shape[-1] == self.codebook_dim
# 将 zhat 进行缩放和偏移
zhat = self._scale_and_shift(zhat)
# 计算索引并转换为整数类型
return (zhat * self._basis).sum(dim=-1).to(int32)
# 定义将索引转换为代码的函数
def indices_to_codes(self, indices: Tensor, project_out=True) -> Tensor:
"""Inverse of `codes_to_indices`."""
# 检查输入是否为图像或视频
is_img_or_video = indices.ndim >= (3 + int(self.keep_num_codebooks_dim))
# 调整索引的形状以增加一个维度
indices = rearrange(indices, "... -> ... 1")
# 计算非中心化的代码
codes_non_centered = (indices // self._basis) % self._levels
# 将非中心化代码进行反缩放和偏移
codes = self._scale_and_shift_inverse(codes_non_centered)
# 如果需要保留代码本维度，则调整代码形状
if self.keep_num_codebooks_dim:
codes = rearrange(codes, "... c d -> ... (c d)")
# 如果需要，进行投影操作
if project_out:
codes = self.project_out(codes)
# 如果是图像或视频，调整代码的形状
if is_img_or_video:
codes = rearrange(codes, "b ... d -> b d ...")
# 返回处理后的代码
return codes
# 定义前向传播函数
@autocast(enabled=False)
def forward(self, z: Tensor) -> Tensor:
"""
einstein notation
b - batch
n - sequence (or flattened spatial dimensions)
d - feature dimension
c - number of codebook dim
"""
# 检查输入是否为图像或视频
is_img_or_video = z.ndim >= 4
# 将图像或视频标准化为 (batch, seq, dimension)
if is_img_or_video:
# 调整 z 的形状
z = rearrange(z, "b d ... -> b ... d")
# 打包 z 以获取适当的维度信息
z, ps = pack_one(z, "b * d")
# 确保 z 的最后一维与期望的维度匹配
assert z.shape[-1] == self.dim, f"expected dimension of {self.dim} but found dimension of {z.shape[-1]}"
# 将输入 z 进行投影
z = self.project_in(z)
# 调整 z 的形状以便于后续处理
z = rearrange(z, "b n (c d) -> b n c d", c=self.num_codebooks)
# 对 z 进行量化，得到代码
codes = self.quantize(z)
# 将代码转换为索引
indices = self.codes_to_indices(codes)
# 调整代码的形状
codes = rearrange(codes, "b n c d -> b n (c d)")
# 对代码进行投影以生成输出
out = self.project_out(codes)
# 重新构造图像或视频的维度
if is_img_or_video:
# 解包输出以恢复原始形状
out = unpack_one(out, ps, "b * d")
# 调整输出的形状
out = rearrange(out, "b ... d -> b d ...")
# 解包索引以恢复原始形状
indices = unpack_one(indices, ps, "b * c")
# 如果不保留代码本维度，则调整索引的形状
if not self.keep_num_codebooks_dim:
indices = rearrange(indices, "... 1 -> ...")
# 返回最终的输出和索引
return out, indices
.\cogvideo-finetune\sat\sgm\modules\autoencoding\regularizers\lookup_free_quantization.py
# 文档字符串：查找自由量化的说明和论文链接
"""
Lookup Free Quantization
Proposed in https://arxiv.org/abs/2310.05737
In the simplest setup, each dimension is quantized into {-1, 1}.
An entropy penalty is used to encourage utilization.
"""
# 从数学库导入对数和向上取整的函数
from math import log2, ceil
# 从 collections 导入命名元组，用于创建简单的类
from collections import namedtuple
# 导入 PyTorch 库及其模块
import torch
from torch import nn, einsum
import torch.nn.functional as F
from torch.nn import Module
from torch.cuda.amp import autocast
# 导入 einops 库，用于张量操作
from einops import rearrange, reduce, pack, unpack
# 常量定义
# 创建一个命名元组，用于存储量化结果及其索引和熵辅助损失
Return = namedtuple("Return", ["quantized", "indices", "entropy_aux_loss"])
# 创建一个命名元组，用于存储损失分解的信息
LossBreakdown = namedtuple("LossBreakdown", ["per_sample_entropy", "batch_entropy", "commitment"])
# 辅助函数
# 检查变量是否存在（非 None）
def exists(v):
return v is not None
# 返回第一个非 None 的参数
def default(*args):
for arg in args:
if exists(arg):
return arg() if callable(arg) else arg
return None
# 将一个张量按模式打包
def pack_one(t, pattern):
return pack([t], pattern)
# 按模式解包张量，返回第一个解包的元素
def unpack_one(t, ps, pattern):
return unpack(t, ps, pattern)[0]
# 熵相关函数
# 计算张量的对数，确保最小值不小于 eps
def log(t, eps=1e-5):
return t.clamp(min=eps).log()
# 计算概率分布的熵
def entropy(prob):
return (-prob * log(prob)).sum(dim=-1)
# 类定义
# 定义 LFQ 类，继承自 PyTorch 的 Module
class LFQ(Module):
# 初始化函数，设置多个参数
def __init__(
self,
*,
dim=None, # 量化维度
codebook_size=None, # 码本大小
entropy_loss_weight=0.1, # 熵损失的权重
commitment_loss_weight=0.25, # 承诺损失的权重
diversity_gamma=1.0, # 多样性控制参数
straight_through_activation=nn.Identity(), # 直通激活函数
num_codebooks=1, # 码本数量
keep_num_codebooks_dim=None, # 保持码本维度
codebook_scale=1.0, # 码本缩放因子，残差 LFQ 每层缩小 2 倍
frac_per_sample_entropy=1.0, # 每个样本熵的比例，若小于 1 则随机使用部分概率
):
# 调用父类构造函数初始化
super().__init__()
# 一些断言验证
# 确保至少指定 dim 或 codebook_size
assert exists(dim) or exists(codebook_size), "either dim or codebook_size must be specified for LFQ"
# 确保 codebook_size 是 2 的幂，若指定了
assert (
not exists(codebook_size) or log2(codebook_size).is_integer()
), f"your codebook size must be a power of 2 for lookup free quantization (suggested {2 ** ceil(log2(codebook_size))})"
# 若未指定 codebook_size，则使用默认值 2 的 dim 次方
codebook_size = default(codebook_size, lambda: 2**dim)
# 计算 codebook 的维度
codebook_dim = int(log2(codebook_size))
# 计算总的 codebook 维度
codebook_dims = codebook_dim * num_codebooks
# 若未指定 dim，则使用 codebook_dims
dim = default(dim, codebook_dims)
# 检查是否存在投影
has_projections = dim != codebook_dims
# 根据是否有投影选择线性层或身份映射
self.project_in = nn.Linear(dim, codebook_dims) if has_projections else nn.Identity()
self.project_out = nn.Linear(codebook_dims, dim) if has_projections else nn.Identity()
# 存储是否有投影的布尔值
self.has_projections = has_projections
# 保存维度和相关参数
self.dim = dim
self.codebook_dim = codebook_dim
self.num_codebooks = num_codebooks
# 处理保持 codebook 维度的默认值
keep_num_codebooks_dim = default(keep_num_codebooks_dim, num_codebooks > 1)
# 确保在多 codebook 时要保持维度
assert not (num_codebooks > 1 and not keep_num_codebooks_dim)
self.keep_num_codebooks_dim = keep_num_codebooks_dim
# 直通激活函数
self.activation = straight_through_activation
# 与熵辅助损失相关的权重
# 确保熵比例在合理范围内
assert 0 < frac_per_sample_entropy <= 1.0
self.frac_per_sample_entropy = frac_per_sample_entropy
# 保存熵损失的权重和其他参数
self.diversity_gamma = diversity_gamma
self.entropy_loss_weight = entropy_loss_weight
# 代码簿缩放因子
self.codebook_scale = codebook_scale
# 承诺损失的权重
self.commitment_loss_weight = commitment_loss_weight
# 用于推理时没有辅助损失的情况
# 注册一个掩码，用于后续计算
self.register_buffer("mask", 2 ** torch.arange(codebook_dim - 1, -1, -1))
# 注册一个零张量，非持久化
self.register_buffer("zero", torch.tensor(0.0), persistent=False)
# 代码的初始化
# 创建所有可能的代码
all_codes = torch.arange(codebook_size)
# 通过掩码生成二进制位的浮点表示
bits = ((all_codes[..., None].int() & self.mask) != 0).float()
# 将二进制位转换为代码簿
codebook = self.bits_to_codes(bits)
# 注册代码簿，非持久化
self.register_buffer("codebook", codebook, persistent=False)
# 将位转换为代码的方法
def bits_to_codes(self, bits):
return bits * self.codebook_scale * 2 - self.codebook_scale
# dtype 属性的定义
@property
def dtype(self):
return self.codebook.dtype
# 将索引转换为代码，返回相应的编码
def indices_to_codes(self, indices, project_out=True):
# 检查索引是否为图像或视频数据，判断维度
is_img_or_video = indices.ndim >= (3 + int(self.keep_num_codebooks_dim))
# 如果不保留代码本维度，将索引重排列为增加一个维度
if not self.keep_num_codebooks_dim:
indices = rearrange(indices, "... -> ... 1")
# 将索引转换为代码，生成-1或1的位
bits = ((indices[..., None].int() & self.mask) != 0).to(self.dtype)
# 将位转换为编码
codes = self.bits_to_codes(bits)
# 将编码重排列为合并维度
codes = rearrange(codes, "... c d -> ... (c d)")
# 判断是否将编码投影回原始维度
if project_out:
codes = self.project_out(codes)
# 将编码重排列回原始形状
if is_img_or_video:
codes = rearrange(codes, "b ... d -> b d ...")
# 返回最终的编码
return codes
# 禁用自动混合精度
@autocast(enabled=False)
def forward(
self,
x,
inv_temperature=100.0,
return_loss_breakdown=False,
mask=None,
.\cogvideo-finetune\sat\sgm\modules\autoencoding\regularizers\quantize.py
# 导入 logging 模块以便进行日志记录
import logging
# 从 abc 模块导入 abstractmethod，用于定义抽象方法
from abc import abstractmethod
# 从 typing 模块导入多种类型提示
from typing import Dict, Iterator, Literal, Optional, Tuple, Union
# 导入 numpy 库并命名为 np
import numpy as np
# 导入 PyTorch 库及其子模块
import torch
import torch.nn as nn
import torch.nn.functional as F
# 从 einops 导入 rearrange 函数，用于重排张量
from einops import rearrange
# 从 torch 导入 einsum 函数，用于张量操作
from torch import einsum
# 从同一包中导入 AbstractRegularizer 类和 measure_perplexity 函数
from .base import AbstractRegularizer, measure_perplexity
# 创建一个 logger 实例，用于当前模块的日志记录
logpy = logging.getLogger(__name__)
# 定义一个抽象量化器类，继承自 AbstractRegularizer
class AbstractQuantizer(AbstractRegularizer):
# 初始化方法
def __init__(self):
# 调用父类的初始化方法
super().__init__()
# 在初始化时定义这些属性
# shape (N,) 表示该张量的形状为一维
self.used: Optional[torch.Tensor] # 定义已使用的张量，可能为 None
self.re_embed: int # 定义重嵌入的整数值
self.unknown_index: Union[Literal["random"], int] # 定义未知索引，可能为随机或整数
# 将输入索引映射到已使用的索引
def remap_to_used(self, inds: torch.Tensor) -> torch.Tensor:
# 确保已定义 used 索引
assert self.used is not None, "You need to define used indices for remap"
ishape = inds.shape # 获取输入索引的形状
assert len(ishape) > 1 # 确保输入维度大于 1
inds = inds.reshape(ishape[0], -1) # 重塑输入索引为二维
used = self.used.to(inds) # 将 used 张量移动到与 inds 相同的设备
match = (inds[:, :, None] == used[None, None, ...]).long() # 计算索引匹配情况
new = match.argmax(-1) # 找到每个匹配的最大索引
unknown = match.sum(2) < 1 # 标记未知索引
# 如果未知索引为随机，则随机生成新的索引
if self.unknown_index == "random":
new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to(device=new.device)
else:
new[unknown] = self.unknown_index # 将未知索引设置为指定的未知索引
return new.reshape(ishape) # 返回重塑后的新索引
# 将输入索引映射回所有索引
def unmap_to_all(self, inds: torch.Tensor) -> torch.Tensor:
# 确保已定义 used 索引
assert self.used is not None, "You need to define used indices for remap"
ishape = inds.shape # 获取输入索引的形状
assert len(ishape) > 1 # 确保输入维度大于 1
inds = inds.reshape(ishape[0], -1) # 重塑输入索引为二维
used = self.used.to(inds) # 将 used 张量移动到与 inds 相同的设备
# 如果重嵌入数量大于已使用数量，则处理额外的令牌
if self.re_embed > self.used.shape[0]: # extra token
inds[inds >= self.used.shape[0]] = 0 # 将超出范围的索引设置为零
back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds) # 根据输入索引收集数据
return back.reshape(ishape) # 返回重塑后的数据
# 定义抽象方法以获取编码表条目
@abstractmethod
def get_codebook_entry(self, indices: torch.Tensor, shape: Optional[Tuple[int, ...]] = None) -> torch.Tensor:
raise NotImplementedError() # 抛出未实现错误
# 获取可训练参数的迭代器
def get_trainable_parameters(self) -> Iterator[torch.nn.Parameter]:
yield from self.parameters() # 生成模型参数
# 定义 Gumbel 量化器类，继承自 AbstractQuantizer
class GumbelQuantizer(AbstractQuantizer):
"""
credit to @karpathy:
https://github.com/karpathy/deep-vector-quantization/blob/main/model.py (thanks!)
Gumbel Softmax trick quantizer
Categorical Reparameterization with Gumbel-Softmax, Jang et al. 2016
https://arxiv.org/abs/1611.01144
"""
# 初始化方法
def __init__(
self,
num_hiddens: int, # 隐藏层单元数
embedding_dim: int, # 嵌入维度
n_embed: int, # 嵌入数量
straight_through: bool = True, # 是否使用直通梯度
kl_weight: float = 5e-4, # KL 散度的权重
temp_init: float = 1.0, # 初始化温度
remap: Optional[str] = None, # 可选的重映射方式
unknown_index: str = "random", # 未知索引的默认值为随机
loss_key: str = "loss/vq", # 损失键的默认值
# 定义一个返回 None 的方法
) -> None:
# 调用父类的构造函数
super().__init__()
# 保存损失的关键字
self.loss_key = loss_key
# 设置嵌入维度
self.embedding_dim = embedding_dim
# 设置嵌入数量
self.n_embed = n_embed
# 设置是否使用直通估计
self.straight_through = straight_through
# 初始化温度参数
self.temperature = temp_init
# 设置 KL 散度权重
self.kl_weight = kl_weight
# 创建一个 2D 卷积层，输入通道为 num_hiddens，输出通道为 n_embed，卷积核大小为 1
self.proj = nn.Conv2d(num_hiddens, n_embed, 1)
# 创建嵌入层，嵌入数量为 n_embed，嵌入维度为 embedding_dim
self.embed = nn.Embedding(n_embed, embedding_dim)
# 保存重映射文件路径
self.remap = remap
# 如果提供了重映射
if self.remap is not None:
# 从重映射文件中加载使用的索引，并将其注册为缓冲区
self.register_buffer("used", torch.tensor(np.load(self.remap)))
# 设置重嵌入数量为使用的索引的数量
self.re_embed = self.used.shape[0]
else:
# 如果未提供重映射，则使用全部嵌入数量
self.used = None
self.re_embed = n_embed
# 如果未知索引设置为 "extra"
if unknown_index == "extra":
# 将未知索引设置为重嵌入数量，并增加一个
self.unknown_index = self.re_embed
self.re_embed = self.re_embed + 1
else:
# 断言未知索引必须为 "random"、"extra" 或整数
assert unknown_index == "random" or isinstance(
unknown_index, int
), "unknown index needs to be 'random', 'extra' or any integer"
# 设置未知索引
self.unknown_index = unknown_index # "random" or "extra" or integer
# 如果提供了重映射，则记录相关信息
if self.remap is not None:
logpy.info(
f"Remapping {self.n_embed} indices to {self.re_embed} indices. "
f"Using {self.unknown_index} for unknown indices."
)
# 定义前向传播方法，接收输入张量 z 和可选的温度参数
def forward(
self, z: torch.Tensor, temp: Optional[float] = None, return_logits: bool = False
) -> Tuple[torch.Tensor, Dict]:
# 在评估模式下强制 hard=True，因为必须进行量化。
# 实际上，始终为真似乎也有效
hard = self.straight_through if self.training else True
# 设置温度，如果未提供，则使用默认值
temp = self.temperature if temp is None else temp
# 初始化输出字典
out_dict = {}
# 通过卷积层计算 logits
logits = self.proj(z)
# 如果提供了重映射
if self.remap is not None:
# 创建与 logits 同样形状的全零张量
full_zeros = torch.zeros_like(logits)
# 仅保留使用的 logits
logits = logits[:, self.used, ...]
# 使用 Gumbel-Softmax 函数生成软 one-hot 编码
soft_one_hot = F.gumbel_softmax(logits, tau=temp, dim=1, hard=hard)
# 如果提供了重映射
if self.remap is not None:
# 将未使用的条目设置为零
full_zeros[:, self.used, ...] = soft_one_hot
soft_one_hot = full_zeros
# 根据软 one-hot 编码和嵌入权重计算量化的 z
z_q = einsum("b n h w, n d -> b d h w", soft_one_hot, self.embed.weight)
# 计算 KL 散度损失
qy = F.softmax(logits, dim=1)
# 计算散度并取平均
diff = self.kl_weight * torch.sum(qy * torch.log(qy * self.n_embed + 1e-10), dim=1).mean()
# 将散度损失存储到输出字典中
out_dict[self.loss_key] = diff
# 计算 soft_one_hot 编码的最大索引
ind = soft_one_hot.argmax(dim=1)
# 如果提供了重映射，将索引转换为使用的索引
if self.remap is not None:
ind = self.remap_to_used(ind)
# 如果需要返回 logits，则将其存储到输出字典中
if return_logits:
out_dict["logits"] = logits
# 返回量化的 z 和输出字典
return z_q, out_dict
# 获取代码本条目的方法，根据给定的索引和形状
def get_codebook_entry(self, indices, shape):
# TODO: 当前形状参数尚不可选
b, h, w, c = shape # 解包形状参数，获取批次、身高、宽度和通道数
# 确保索引的总数与给定的形状匹配
assert b * h * w == indices.shape[0]
# 重排列索引，将其形状调整为 (b, h, w)
indices = rearrange(indices, "(b h w) -> b h w", b=b, h=h, w=w)
# 如果存在重映射，则将索引映射到所有可能的值
if self.remap is not None:
indices = self.unmap_to_all(indices)
# 将索引转换为独热编码，调整维度顺序并转换为浮点数
one_hot = F.one_hot(indices, num_classes=self.n_embed).permute(0, 3, 1, 2).float()
# 通过爱因斯坦求和约定计算最终的量化表示 z_q
z_q = einsum("b n h w, n d -> b d h w", one_hot, self.embed.weight)
# 返回量化后的表示
return z_q
# 定义向量量化类，继承自抽象量化器
class VectorQuantizer(AbstractQuantizer):
"""
____________________________________________
VQ-VAE 的离散化瓶颈部分。
输入:
- n_e : 嵌入的数量
- e_dim : 嵌入的维度
- beta : 在损失项中使用的承诺成本,
beta * ||z_e(x)-sg[e]||^2
_____________________________________________
"""
# 初始化方法，定义类的参数
def __init__(
self,
n_e: int, # 嵌入的数量
e_dim: int, # 嵌入的维度
beta: float = 0.25, # 默认承诺成本
remap: Optional[str] = None, # 可选的重映射文件路径
unknown_index: str = "random", # 未知索引的处理方式
sane_index_shape: bool = False, # 是否保持合理的索引形状
log_perplexity: bool = False, # 是否记录困惑度
embedding_weight_norm: bool = False, # 是否使用权重归一化
loss_key: str = "loss/vq", # 损失的键
):
# 调用父类的初始化方法
super().__init__()
# 保存嵌入数量
self.n_e = n_e
# 保存嵌入维度
self.e_dim = e_dim
# 保存承诺成本
self.beta = beta
# 保存损失键
self.loss_key = loss_key
# 如果不使用权重归一化
if not embedding_weight_norm:
# 初始化嵌入层，权重范围均匀分布
self.embedding = nn.Embedding(self.n_e, self.e_dim)
self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)
else:
# 使用权重归一化的嵌入层
self.embedding = torch.nn.utils.weight_norm(nn.Embedding(self.n_e, self.e_dim), dim=1)
# 保存重映射参数
self.remap = remap
# 如果指定了重映射
if self.remap is not None:
# 从重映射文件中加载已使用的索引
self.register_buffer("used", torch.tensor(np.load(self.remap)))
# 设置重新嵌入的数量
self.re_embed = self.used.shape[0]
else:
# 否则未使用的索引为 None
self.used = None
# 重新嵌入的数量为 n_e
self.re_embed = n_e
# 如果未知索引是 "extra"
if unknown_index == "extra":
# 设置未知索引为重新嵌入的数量
self.unknown_index = self.re_embed
# 重新嵌入的数量加一
self.re_embed = self.re_embed + 1
else:
# 确保未知索引是 "random"、"extra" 或整数
assert unknown_index == "random" or isinstance(
unknown_index, int
), "unknown index needs to be 'random', 'extra' or any integer"
# 保存未知索引的值
self.unknown_index = unknown_index # "random" 或 "extra" 或整数
# 如果指定了重映射，记录信息
if self.remap is not None:
logpy.info(
f"Remapping {self.n_e} indices to {self.re_embed} indices. "
f"Using {self.unknown_index} for unknown indices."
)
# 保存是否保持合理的索引形状的标志
self.sane_index_shape = sane_index_shape
# 保存是否记录困惑度的标志
self.log_perplexity = log_perplexity
# 前向传播方法，定义输入的处理
def forward(
self,
z: torch.Tensor, # 输入张量
) -> Tuple[torch.Tensor, Dict]: # 定义返回类型为元组，包含一个张量和一个字典
do_reshape = z.ndim == 4 # 检查 z 的维度是否为 4，决定是否需要重塑
if do_reshape: # 如果 z 是 4 维的
# reshape z -> (batch, height, width, channel) and flatten # 重塑 z 的维度为 (batch, height, width, channel) 并扁平化
z = rearrange(z, "b c h w -> b h w c").contiguous() # 重新排列 z 的维度，并保证内存连续性
else: # 如果 z 不是 4 维的
assert z.ndim < 4, "No reshaping strategy for inputs > 4 dimensions defined" # 断言 z 的维度小于 4
z = z.contiguous() # 确保 z 的内存是连续的
z_flattened = z.view(-1, self.e_dim) # 将 z 重塑为 (batch_size, e_dim) 的形状
# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z # 计算 z 到嵌入 e_j 的距离
d = ( # 计算每个 z 到嵌入的距离
torch.sum(z_flattened**2, dim=1, keepdim=True) # 计算 z_flattened 的平方和
+ torch.sum(self.embedding.weight**2, dim=1) # 计算嵌入权重的平方和
- 2 * torch.einsum("bd,dn->bn", z_flattened, rearrange(self.embedding.weight, "n d -> d n")) # 计算 z_flattened 与嵌入的内积
)
min_encoding_indices = torch.argmin(d, dim=1) # 找到每个 z 到嵌入距离的最小索引
z_q = self.embedding(min_encoding_indices).view(z.shape) # 根据最小索引获取量化的嵌入，并重塑为原始 z 的形状
loss_dict = {} # 初始化损失字典
if self.log_perplexity: # 如果需要记录困惑度
perplexity, cluster_usage = measure_perplexity(min_encoding_indices.detach(), self.n_e) # 计算困惑度和集群使用情况
loss_dict.update({"perplexity": perplexity, "cluster_usage": cluster_usage}) # 更新损失字典
# compute loss for embedding # 计算嵌入的损失
loss = self.beta * torch.mean((z_q.detach() - z) ** 2) + torch.mean((z_q - z.detach()) ** 2) # 计算损失值
loss_dict[self.loss_key] = loss # 将损失添加到损失字典中
# preserve gradients # 保留梯度
z_q = z + (z_q - z).detach() # 将量化的 z_q 与原始 z 结合，保留梯度
# reshape back to match original input shape # 重新调整形状以匹配原始输入形状
if do_reshape: # 如果之前进行了重塑
z_q = rearrange(z_q, "b h w c -> b c h w").contiguous() # 将 z_q 的维度调整回 (batch, channel, height, width)
if self.remap is not None: # 如果需要重映射
min_encoding_indices = min_encoding_indices.reshape(z.shape[0], -1) # 添加批次维度
min_encoding_indices = self.remap_to_used(min_encoding_indices) # 对最小编码索引进行重映射
min_encoding_indices = min_encoding_indices.reshape(-1, 1) # 扁平化为一维
if self.sane_index_shape: # 如果索引形状正常
if do_reshape: # 如果之前进行了重塑
min_encoding_indices = min_encoding_indices.reshape(z_q.shape[0], z_q.shape[2], z_q.shape[3]) # 将索引重塑为与 z_q 形状相匹配
else: # 如果没有重塑
min_encoding_indices = rearrange(min_encoding_indices, "(b s) 1 -> b s", b=z_q.shape[0]) # 重新排列为 (batch, size)
loss_dict["min_encoding_indices"] = min_encoding_indices # 将最小编码索引添加到损失字典中
return z_q, loss_dict # 返回量化的 z_q 和损失字典
def get_codebook_entry(self, indices: torch.Tensor, shape: Optional[Tuple[int, ...]] = None) -> torch.Tensor: # 定义方法获取代码本条目
# shape specifying (batch, height, width, channel) # shape 指定 (batch, height, width, channel)
if self.remap is not None: # 如果需要重映射
assert shape is not None, "Need to give shape for remap" # 断言必须提供形状以进行重映射
indices = indices.reshape(shape[0], -1) # 添加批次维度
indices = self.unmap_to_all(indices) # 对索引进行反向映射
indices = indices.reshape(-1) # 再次扁平化
# get quantized latent vectors # 获取量化的潜在向量
z_q = self.embedding(indices) # 根据索引获取嵌入
if shape is not None: # 如果提供了形状
z_q = z_q.view(shape) # 将 z_q 重塑为指定的形状
# reshape back to match original input shape # 重新调整形状以匹配原始输入形状
z_q = z_q.permute(0, 3, 1, 2).contiguous() # 调整维度顺序并确保内存连续
return z_q # 返回量化后的 z_q
# 定义一个名为 EmbeddingEMA 的神经网络模块，继承自 nn.Module
class EmbeddingEMA(nn.Module):
# 初始化函数，接收令牌数量、码本维度、衰减因子和小常数作为参数
def __init__(self, num_tokens, codebook_dim, decay=0.99, eps=1e-5):
# 调用父类的初始化方法
super().__init__()
# 设置衰减因子
self.decay = decay
# 设置小常数以避免除零
self.eps = eps
# 生成一个随机的权重矩阵，形状为 (num_tokens, codebook_dim)
weight = torch.randn(num_tokens, codebook_dim)
# 将权重定义为不可训练的参数
self.weight = nn.Parameter(weight, requires_grad=False)
# 初始化集群大小为零，定义为不可训练的参数
self.cluster_size = nn.Parameter(torch.zeros(num_tokens), requires_grad=False)
# 复制权重并将其定义为不可训练的参数，用于存储嵌入平均值
self.embed_avg = nn.Parameter(weight.clone(), requires_grad=False)
# 设置更新标志为真
self.update = True
# 前向传播函数，接收嵌入 ID 并返回对应的嵌入向量
def forward(self, embed_id):
return F.embedding(embed_id, self.weight)
# 更新集群大小的指数移动平均
def cluster_size_ema_update(self, new_cluster_size):
# 按衰减因子更新当前集群大小
self.cluster_size.data.mul_(self.decay).add_(new_cluster_size, alpha=1 - self.decay)
# 更新嵌入平均值的指数移动平均
def embed_avg_ema_update(self, new_embed_avg):
# 按衰减因子更新当前嵌入平均值
self.embed_avg.data.mul_(self.decay).add_(new_embed_avg, alpha=1 - self.decay)
# 更新权重，基于平滑的集群大小
def weight_update(self, num_tokens):
# 计算集群大小的总和
n = self.cluster_size.sum()
# 计算平滑的集群大小，以避免除零错误
smoothed_cluster_size = (self.cluster_size + self.eps) / (n + num_tokens * self.eps) * n
# 用平滑的集群大小对嵌入平均值进行归一化
embed_normalized = self.embed_avg / smoothed_cluster_size.unsqueeze(1)
# 用归一化的嵌入更新权重
self.weight.data.copy_(embed_normalized)
# 定义一个名为 EMAVectorQuantizer 的抽象量化器类，继承自 AbstractQuantizer
class EMAVectorQuantizer(AbstractQuantizer):
# 初始化函数，接收嵌入数量、嵌入维度、β、衰减因子、小常数和其他参数
def __init__(
self,
n_embed: int,
embedding_dim: int,
beta: float,
decay: float = 0.99,
eps: float = 1e-5,
remap: Optional[str] = None,
unknown_index: str = "random",
loss_key: str = "loss/vq",
):
# 调用父类的初始化方法
super().__init__()
# 设置码本维度
self.codebook_dim = embedding_dim
# 设置令牌数量
self.num_tokens = n_embed
# 设置 β 值
self.beta = beta
# 设置损失键
self.loss_key = loss_key
# 初始化嵌入 EMA 模块
self.embedding = EmbeddingEMA(self.num_tokens, self.codebook_dim, decay, eps)
# 处理重映射参数
self.remap = remap
if self.remap is not None:
# 如果提供了重映射路径，则加载并注册重映射后的索引
self.register_buffer("used", torch.tensor(np.load(self.remap)))
# 重新嵌入的数量为重映射后索引的形状
self.re_embed = self.used.shape[0]
else:
# 如果没有重映射，则用原令牌数量初始化
self.used = None
self.re_embed = n_embed
# 处理未知索引
if unknown_index == "extra":
# 如果未知索引为 "extra"，则更新重新嵌入数量
self.unknown_index = self.re_embed
self.re_embed = self.re_embed + 1
else:
# 确保未知索引为有效类型
assert unknown_index == "random" or isinstance(
unknown_index, int
), "unknown index needs to be 'random', 'extra' or any integer"
# 设置未知索引为提供的值
self.unknown_index = unknown_index # "random" or "extra" or integer
# 如果存在重映射，则记录重映射信息
if self.remap is not None:
logpy.info(
f"Remapping {self.n_embed} indices to {self.re_embed} indices. "
f"Using {self.unknown_index} for unknown indices."
)
# 定义前向传播函数，接受一个张量 z，返回量化后的张量和字典
def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, Dict]:
# 将 z 的形状调整为 (batch, height, width, channel) 并扁平化
# z, 'b c h w -> b h w c'
z = rearrange(z, "b c h w -> b h w c") # 调整 z 的维度顺序
z_flattened = z.reshape(-1, self.codebook_dim) # 将 z 扁平化为二维张量
# 计算 z 与嵌入 e_j 的距离 (z - e)^2 = z^2 + e^2 - 2 e * z
d = (
z_flattened.pow(2).sum(dim=1, keepdim=True) # 计算 z 的平方和
+ self.embedding.weight.pow(2).sum(dim=1) # 计算嵌入的平方和
- 2 * torch.einsum("bd,nd->bn", z_flattened, self.embedding.weight) # 计算 z 和嵌入的点积
) # 'n d -> d n'
# 找到每个 z 的最小距离对应的编码索引
encoding_indices = torch.argmin(d, dim=1)
# 根据编码索引获取量化后的 z，并调整形状以匹配原始 z
z_q = self.embedding(encoding_indices).view(z.shape)
# 对编码进行独热编码，转换为与 z 相同的数据类型
encodings = F.one_hot(encoding_indices, self.num_tokens).type(z.dtype)
# 计算编码的平均概率
avg_probs = torch.mean(encodings, dim=0)
# 计算困惑度
perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
# 如果处于训练状态且允许更新嵌入
if self.training and self.embedding.update:
# 更新 EMA 聚类大小
encodings_sum = encodings.sum(0)
self.embedding.cluster_size_ema_update(encodings_sum) # 更新聚类大小的 EMA
# 更新 EMA 嵌入平均值
embed_sum = encodings.transpose(0, 1) @ z_flattened # 计算加权和
self.embedding.embed_avg_ema_update(embed_sum) # 更新嵌入平均值的 EMA
# 规范化嵌入平均值并更新权重
self.embedding.weight_update(self.num_tokens)
# 计算嵌入的损失
loss = self.beta * F.mse_loss(z_q.detach(), z) # 计算量化 z 与原 z 的均方误差损失
# 保留梯度
z_q = z + (z_q - z).detach() # 使用 z 的值加上 z_q 的变化，但不计算梯度
# 将 z_q 的形状调整回原始输入的形状
# z_q, 'b h w c -> b c h w'
z_q = rearrange(z_q, "b h w c -> b c h w") # 恢复到原始的维度顺序
# 创建一个字典以返回损失和其他信息
out_dict = {
self.loss_key: loss, # 将损失放入字典
"encodings": encodings, # 包含独热编码
"encoding_indices": encoding_indices, # 包含编码索引
"perplexity": perplexity, # 包含困惑度
}
# 返回量化后的 z 和输出字典
return z_q, out_dict
# 定义一个带有输入投影的向量量化类，继承自 VectorQuantizer
class VectorQuantizerWithInputProjection(VectorQuantizer):
# 初始化方法，接受输入维度、编码数量、码本维度等参数
def __init__(
self,
input_dim: int, # 输入数据的维度
n_codes: int, # 编码数量
codebook_dim: int, # 码本的维度
beta: float = 1.0, # 调整项的超参数，默认值为1.0
output_dim: Optional[int] = None, # 输出维度，可选
**kwargs, # 其他关键字参数
):
# 调用父类的初始化方法
super().__init__(n_codes, codebook_dim, beta, **kwargs)
# 创建输入投影层，将输入维度映射到码本维度
self.proj_in = nn.Linear(input_dim, codebook_dim)
# 设置输出维度属性
self.output_dim = output_dim
# 如果指定了输出维度，则创建输出投影层
if output_dim is not None:
self.proj_out = nn.Linear(codebook_dim, output_dim)
else:
# 如果没有指定输出维度，则使用恒等映射
self.proj_out = nn.Identity()
# 前向传播方法，接受输入张量并返回量化结果和损失字典
def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, Dict]:
rearr = False # 初始化重排列标志
in_shape = z.shape # 获取输入张量的形状
# 如果输入张量的维度大于3，则进行重排列
if z.ndim > 3:
rearr = self.output_dim is not None # 检查是否需要重排列
# 将输入张量从 (batch, channels, ...) 转换为 (batch, ..., channels)
z = rearrange(z, "b c ... -> b (...) c")
# 将输入张量投影到码本维度
z = self.proj_in(z)
# 调用父类的前向方法进行量化，获得量化结果和损失字典
z_q, loss_dict = super().forward(z)
# 将量化结果通过输出投影层
z_q = self.proj_out(z_q)
# 如果需要重排列，根据输入形状调整输出张量
if rearr:
# 如果输入维度为4，重排列为 (batch, channels, height, width)
if len(in_shape) == 4:
z_q = rearrange(z_q, "b (h w) c -> b c h w ", w=in_shape[-1])
# 如果输入维度为5，重排列为 (batch, channels, time, height, width)
elif len(in_shape) == 5:
z_q = rearrange(z_q, "b (t h w) c -> b c t h w ", w=in_shape[-1], h=in_shape[-2])
else:
# 如果输入维度不支持重排列，则抛出异常
raise NotImplementedError(f"rearranging not available for {len(in_shape)}-dimensional input.")
# 返回量化结果和损失字典
return z_q, loss_dict
.\cogvideo-finetune\sat\sgm\modules\autoencoding\regularizers\__init__.py
# 导入抽象方法装饰器和类型注解
from abc import abstractmethod
# 导入任意类型和元组类型注解
from typing import Any, Tuple
# 导入 PyTorch 相关模块
import torch
import torch.nn as nn
import torch.nn.functional as F
# 从自定义模块导入对角高斯分布
from ....modules.distributions.distributions import DiagonalGaussianDistribution
# 从基类模块导入抽象正则化器
from .base import AbstractRegularizer
# 定义对角高斯正则化器类，继承自抽象正则化器
class DiagonalGaussianRegularizer(AbstractRegularizer):
# 初始化方法，接收一个布尔值参数，默认值为 True
def __init__(self, sample: bool = True):
# 调用父类的初始化方法
super().__init__()
# 设置实例属性 sample
self.sample = sample
# 定义获取可训练参数的方法，返回任意类型
def get_trainable_parameters(self) -> Any:
# 生成一个空的可迭代对象
yield from ()
# 定义前向传播方法，接收一个张量，返回一个元组
def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, dict]:
# 创建一个空字典，用于存储日志信息
log = dict()
# 创建一个对角高斯分布实例，基于输入张量 z
posterior = DiagonalGaussianDistribution(z)
# 如果 sample 为 True，进行采样
if self.sample:
z = posterior.sample()
# 否则，使用模式值
else:
z = posterior.mode()
# 计算 KL 散度损失
kl_loss = posterior.kl()
# 对 KL 损失求和并平均
kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
# 将 KL 损失添加到日志字典中
log["kl_loss"] = kl_loss
# 返回处理后的张量和日志字典
return z, log
.\cogvideo-finetune\sat\sgm\modules\autoencoding\temporal_ae.py
# 从 typing 模块导入 Callable, Iterable, Union 类型注解
from typing import Callable, Iterable, Union
# 导入 PyTorch 库
import torch
# 从 einops 导入 rearrange 和 repeat 函数，用于张量重排和重复
from einops import rearrange, repeat
# 从自定义模块中导入所需的类和变量
from sgm.modules.diffusionmodules.model import (
# 检查 XFORMERS 库是否可用
XFORMERS_IS_AVAILABLE,
# 导入注意力块、解码器和其他模块
AttnBlock,
Decoder,
MemoryEfficientAttnBlock,
ResnetBlock,
)
# 从 openaimodel 模块导入 ResBlock 和时间步嵌入函数
from sgm.modules.diffusionmodules.openaimodel import ResBlock, timestep_embedding
# 从 video_attention 模块导入视频变换块
from sgm.modules.video_attention import VideoTransformerBlock
# 从 util 模块导入 partialclass 函数
from sgm.util import partialclass
# 定义一个新的类 VideoResBlock，继承自 ResnetBlock
class VideoResBlock(ResnetBlock):
# 初始化函数，接收多个参数
def __init__(
self,
out_channels, # 输出通道数
*args, # 额外参数
dropout=0.0, # dropout 概率
video_kernel_size=3, # 视频卷积核大小
alpha=0.0, # 混合因子
merge_strategy="learned", # 合并策略
**kwargs, # 其他关键字参数
):
# 调用父类构造函数进行初始化
super().__init__(out_channels=out_channels, dropout=dropout, *args, **kwargs)
# 如果未指定 video_kernel_size，则默认设置
if video_kernel_size is None:
video_kernel_size = [3, 1, 1]
# 创建时间堆栈，使用 ResBlock
self.time_stack = ResBlock(
channels=out_channels, # 通道数
emb_channels=0, # 嵌入通道数
dropout=dropout, # dropout 概率
dims=3, # 数据维度
use_scale_shift_norm=False, # 是否使用缩放平移归一化
use_conv=False, # 是否使用卷积
up=False, # 是否向上采样
down=False, # 是否向下采样
kernel_size=video_kernel_size, # 卷积核大小
use_checkpoint=False, # 是否使用检查点
skip_t_emb=True, # 是否跳过时间嵌入
)
# 设置合并策略
self.merge_strategy = merge_strategy
# 根据合并策略注册混合因子
if self.merge_strategy == "fixed":
self.register_buffer("mix_factor", torch.Tensor([alpha])) # 固定混合因子
elif self.merge_strategy == "learned":
self.register_parameter("mix_factor", torch.nn.Parameter(torch.Tensor([alpha]))) # 学习混合因子
else:
raise ValueError(f"unknown merge strategy {self.merge_strategy}") # 抛出未知合并策略错误
# 获取 alpha 值的函数
def get_alpha(self, bs):
# 根据合并策略返回混合因子
if self.merge_strategy == "fixed":
return self.mix_factor # 固定策略
elif self.merge_strategy == "learned":
return torch.sigmoid(self.mix_factor) # 学习策略
else:
raise NotImplementedError() # 抛出未实现错误
# 前向传播函数
def forward(self, x, temb, skip_video=False, timesteps=None):
# 如果未提供时间步，则使用类中的 timesteps
if timesteps is None:
timesteps = self.timesteps
# 获取输入张量的形状
b, c, h, w = x.shape
# 调用父类的前向传播方法
x = super().forward(x, temb)
# 如果不跳过视频处理
if not skip_video:
# 重排张量，将其调整为视频格式
x_mix = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps)
# 重排当前张量
x = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps)
# 通过时间堆栈进行处理
x = self.time_stack(x, temb)
# 获取 alpha 值
alpha = self.get_alpha(bs=b // timesteps)
# 按比例混合两个张量
x = alpha * x + (1.0 - alpha) * x_mix
# 再次重排张量
x = rearrange(x, "b c t h w -> (b t) c h w")
return x # 返回处理后的张量
# 定义一个新的类 AE3DConv，继承自 torch.nn.Conv2d
class AE3DConv(torch.nn.Conv2d):
# 初始化方法，设置输入和输出通道及卷积核大小等参数
def __init__(self, in_channels, out_channels, video_kernel_size=3, *args, **kwargs):
# 调用父类的初始化方法，传递输入和输出通道及其他参数
super().__init__(in_channels, out_channels, *args, **kwargs)
# 检查 video_kernel_size 是否为可迭代对象
if isinstance(video_kernel_size, Iterable):
# 如果是可迭代对象，计算每个核的填充大小
padding = [int(k // 2) for k in video_kernel_size]
else:
# 否则，计算单个核的填充大小
padding = int(video_kernel_size // 2)
# 创建一个 3D 卷积层，用于处理视频数据
self.time_mix_conv = torch.nn.Conv3d(
in_channels=out_channels, # 输入通道数
out_channels=out_channels, # 输出通道数
kernel_size=video_kernel_size, # 卷积核大小
padding=padding, # 填充大小
)
# 前向传播方法，处理输入数据并返回输出
def forward(self, input, timesteps, skip_video=False):
# 调用父类的前向传播方法，处理输入数据
x = super().forward(input)
# 如果跳过视频处理，直接返回处理后的数据
if skip_video:
return x
# 调整张量形状，以适应 3D 卷积的输入要求
x = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps)
# 通过 3D 卷积层处理数据
x = self.time_mix_conv(x)
# 调整输出张量的形状，返回到原始格式
return rearrange(x, "b c t h w -> (b t) c h w")
# 定义一个视频块类，继承自注意力块基类
class VideoBlock(AttnBlock):
# 初始化函数，接收输入通道数、混合因子和合并策略
def __init__(self, in_channels: int, alpha: float = 0, merge_strategy: str = "learned"):
# 调用基类初始化函数
super().__init__(in_channels)
# 创建视频转换块，使用单头注意力机制
self.time_mix_block = VideoTransformerBlock(
dim=in_channels,
n_heads=1,
d_head=in_channels,
checkpoint=False,
ff_in=True,
attn_mode="softmax",
)
# 计算时间嵌入维度
time_embed_dim = self.in_channels * 4
# 构建视频时间嵌入的神经网络结构
self.video_time_embed = torch.nn.Sequential(
torch.nn.Linear(self.in_channels, time_embed_dim),
torch.nn.SiLU(),
torch.nn.Linear(time_embed_dim, self.in_channels),
)
# 设置合并策略
self.merge_strategy = merge_strategy
# 根据合并策略注册混合因子为缓冲区
if self.merge_strategy == "fixed":
self.register_buffer("mix_factor", torch.Tensor([alpha]))
# 注册混合因子为可学习参数
elif self.merge_strategy == "learned":
self.register_parameter("mix_factor", torch.nn.Parameter(torch.Tensor([alpha])))
# 如果合并策略未知，抛出错误
else:
raise ValueError(f"unknown merge strategy {self.merge_strategy}")
# 前向传播函数，接收输入、时间步和跳过视频标志
def forward(self, x, timesteps, skip_video=False):
# 如果跳过视频，调用基类的前向传播
if skip_video:
return super().forward(x)
# 保存输入数据
x_in = x
# 进行注意力计算
x = self.attention(x)
# 获取输出的高度和宽度
h, w = x.shape[2:]
# 重新排列数据形状
x = rearrange(x, "b c h w -> b (h w) c")
# 初始化混合输入
x_mix = x
# 创建时间步的序列
num_frames = torch.arange(timesteps, device=x.device)
# 重复时间步以匹配批大小
num_frames = repeat(num_frames, "t -> b t", b=x.shape[0] // timesteps)
# 重新排列时间步
num_frames = rearrange(num_frames, "b t -> (b t)")
# 生成时间嵌入
t_emb = timestep_embedding(num_frames, self.in_channels, repeat_only=False)
# 计算时间嵌入的输出
emb = self.video_time_embed(t_emb) # b, n_channels
# 增加一个维度
emb = emb[:, None, :]
# 将时间嵌入与输入混合
x_mix = x_mix + emb
# 获取当前的混合因子
alpha = self.get_alpha()
# 进行时间混合块计算
x_mix = self.time_mix_block(x_mix, timesteps=timesteps)
# 根据混合因子合并输入和混合输出
x = alpha * x + (1.0 - alpha) * x_mix # alpha merge
# 重新排列输出形状
x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w)
# 投影到输出空间
x = self.proj_out(x)
# 返回输入与输出的和
return x_in + x
# 获取当前的混合因子
def get_alpha(
self,
):
# 如果合并策略是固定，返回固定的混合因子
if self.merge_strategy == "fixed":
return self.mix_factor
# 如果合并策略是学习的，返回经过 sigmoid 函数处理的混合因子
elif self.merge_strategy == "learned":
return torch.sigmoid(self.mix_factor)
# 如果合并策略未知，抛出错误
else:
raise NotImplementedError(f"unknown merge strategy {self.merge_strategy}")
# 定义一个内存高效的视频块类，继承自内存高效注意力块
class MemoryEfficientVideoBlock(MemoryEfficientAttnBlock):
# 初始化类，设置输入通道、混合因子和合并策略
def __init__(self, in_channels: int, alpha: float = 0, merge_strategy: str = "learned"):
# 调用父类构造函数，传递输入通道数
super().__init__(in_channels)
# 创建视频变换块，设置相关参数
self.time_mix_block = VideoTransformerBlock(
dim=in_channels,
n_heads=1,
d_head=in_channels,
checkpoint=False,
ff_in=True,
attn_mode="softmax-xformers",
)
# 计算时间嵌入维度
time_embed_dim = self.in_channels * 4
# 定义时间嵌入序列，包含两层线性变换和SiLU激活
self.video_time_embed = torch.nn.Sequential(
torch.nn.Linear(self.in_channels, time_embed_dim),
torch.nn.SiLU(),
torch.nn.Linear(time_embed_dim, self.in_channels),
)
# 保存合并策略
self.merge_strategy = merge_strategy
# 如果合并策略是固定，则注册混合因子为缓冲区
if self.merge_strategy == "fixed":
self.register_buffer("mix_factor", torch.Tensor([alpha]))
# 如果合并策略是学习的，则注册混合因子为可学习参数
elif self.merge_strategy == "learned":
self.register_parameter("mix_factor", torch.nn.Parameter(torch.Tensor([alpha])))
# 否则抛出错误
else:
raise ValueError(f"unknown merge strategy {self.merge_strategy}")
# 前向传播函数
def forward(self, x, timesteps, skip_time_block=False):
# 如果跳过时间块，调用父类的前向传播
if skip_time_block:
return super().forward(x)
# 保存输入数据
x_in = x
# 应用注意力机制
x = self.attention(x)
# 获取输出的高度和宽度
h, w = x.shape[2:]
# 重排张量以便于处理
x = rearrange(x, "b c h w -> b (h w) c")
# 初始化混合输入
x_mix = x
# 创建时间帧的张量
num_frames = torch.arange(timesteps, device=x.device)
# 重复时间帧以匹配批次大小
num_frames = repeat(num_frames, "t -> b t", b=x.shape[0] // timesteps)
# 重排张量
num_frames = rearrange(num_frames, "b t -> (b t)")
# 获取时间嵌入
t_emb = timestep_embedding(num_frames, self.in_channels, repeat_only=False)
# 应用视频时间嵌入
emb = self.video_time_embed(t_emb) # b, n_channels
# 在第二维插入新的维度
emb = emb[:, None, :]
# 将嵌入加到混合输入
x_mix = x_mix + emb
# 获取混合因子
alpha = self.get_alpha()
# 应用时间混合块
x_mix = self.time_mix_block(x_mix, timesteps=timesteps)
# 根据alpha进行混合
x = alpha * x + (1.0 - alpha) * x_mix # alpha merge
# 重新排列张量以恢复原始维度
x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w)
# 应用输出投影
x = self.proj_out(x)
# 返回输入与输出的和
return x_in + x
# 获取混合因子的函数
def get_alpha(
self,
):
# 如果合并策略是固定，则返回混合因子
if self.merge_strategy == "fixed":
return self.mix_factor
# 如果合并策略是学习的，则返回经过sigmoid处理的混合因子
elif self.merge_strategy == "learned":
return torch.sigmoid(self.mix_factor)
# 否则抛出未实现错误
else:
raise NotImplementedError(f"unknown merge strategy {self.merge_strategy}")
# 创建时空注意力机制的函数，接受多个参数以配置注意力类型和其他设置
def make_time_attn(
in_channels, # 输入通道数
attn_type="vanilla", # 注意力类型，默认为'vanilla'
attn_kwargs=None, # 额外的注意力参数，默认为None
alpha: float = 0, # 参数alpha，默认为0
merge_strategy: str = "learned", # 合并策略，默认为'learned'
):
# 检查注意力类型是否在支持的选项中
assert attn_type in [
"vanilla",
"vanilla-xformers",
], f"attn_type {attn_type} not supported for spatio-temporal attention"
# 打印当前创建的注意力类型及其输入通道数
print(f"making spatial and temporal attention of type '{attn_type}' with {in_channels} in_channels")
# 如果不支持xformers，且当前注意力类型为'vanilla-xformers'，则回退到'vanilla'
if not XFORMERS_IS_AVAILABLE and attn_type == "vanilla-xformers":
print(
f"Attention mode '{attn_type}' is not available. Falling back to vanilla attention. "
f"This is not a problem in Pytorch >= 2.0. FYI, you are running with PyTorch version {torch.__version__}"
)
attn_type = "vanilla"
# 如果注意力类型为'vanilla'，则返回部分类VideoBlock
if attn_type == "vanilla":
assert attn_kwargs is None # 确保没有提供额外的参数
return partialclass(VideoBlock, in_channels, alpha=alpha, merge_strategy=merge_strategy)
# 如果注意力类型为'vanilla-xformers'，则返回部分类MemoryEfficientVideoBlock
elif attn_type == "vanilla-xformers":
print(f"building MemoryEfficientAttnBlock with {in_channels} in_channels...")
return partialclass(
MemoryEfficientVideoBlock,
in_channels,
alpha=alpha,
merge_strategy=merge_strategy,
)
else:
return NotImplementedError() # 如果不支持的类型，返回未实现错误
# 自定义的卷积层包装器，继承自torch.nn.Conv2d
class Conv2DWrapper(torch.nn.Conv2d):
# 前向传播方法，调用父类的前向传播
def forward(self, input: torch.Tensor, **kwargs) -> torch.Tensor:
return super().forward(input)
# 视频解码器类，继承自Decoder
class VideoDecoder(Decoder):
available_time_modes = ["all", "conv-only", "attn-only"] # 可用的时间模式列表
# 初始化方法，设置视频解码器的各种参数
def __init__(
self,
*args,
video_kernel_size: Union[int, list] = 3, # 视频卷积核大小，默认为3
alpha: float = 0.0, # alpha参数，默认为0.0
merge_strategy: str = "learned", # 合并策略，默认为'learned'
time_mode: str = "conv-only", # 时间模式，默认为'conv-only'
**kwargs,
):
self.video_kernel_size = video_kernel_size # 设置视频卷积核大小
self.alpha = alpha # 设置alpha参数
self.merge_strategy = merge_strategy # 设置合并策略
self.time_mode = time_mode # 设置时间模式
# 确保时间模式在可用选项内
assert (
self.time_mode in self.available_time_modes
), f"time_mode parameter has to be in {self.available_time_modes}"
super().__init__(*args, **kwargs) # 调用父类的初始化方法
# 获取最后一层的权重，支持跳过时间混合的选项
def get_last_layer(self, skip_time_mix=False, **kwargs):
# 如果时间模式为'attn-only'，则抛出未实现错误
if self.time_mode == "attn-only":
raise NotImplementedError("TODO")
else:
# 返回适当的权重，基于跳过时间混合的选项
return self.conv_out.time_mix_conv.weight if not skip_time_mix else self.conv_out.weight
# 创建注意力机制的方法
def _make_attn(self) -> Callable:
# 根据时间模式返回适当的部分类
if self.time_mode not in ["conv-only", "only-last-conv"]:
return partialclass(
make_time_attn,
alpha=self.alpha,
merge_strategy=self.merge_strategy,
)
else:
return super()._make_attn() # 否则调用父类的方法
# 创建卷积层的方法
def _make_conv(self) -> Callable:
# 根据时间模式返回适当的部分类或卷积包装器
if self.time_mode != "attn-only":
return partialclass(AE3DConv, video_kernel_size=self.video_kernel_size)
else:
return Conv2DWrapper # 返回卷积包装器
# 定义一个私有方法，用于创建残差块，返回一个可调用对象
def _make_resblock(self) -> Callable:
# 检查当前的时间模式是否不在指定的两种模式中
if self.time_mode not in ["attn-only", "only-last-conv"]:
# 返回一个部分应用的类，用于创建 VideoResBlock 实例
return partialclass(
VideoResBlock,
video_kernel_size=self.video_kernel_size,
alpha=self.alpha,
merge_strategy=self.merge_strategy,
)
else:
# 否则，调用父类的方法以创建残差块
return super()._make_resblock()
.\cogvideo-finetune\sat\sgm\modules\autoencoding\vqvae\movq_dec_3d.py
# pytorch_diffusion + derived encoder decoder
# 导入所需的库
import math
import torch
import torch.nn as nn
import numpy as np
from einops import rearrange
from .movq_enc_3d import CausalConv3d, Upsample3D, DownSample3D
# 将输入转换为元组，确保其长度为指定值
def cast_tuple(t, length=1):
return t if isinstance(t, tuple) else ((t,) * length)
# 检查一个数是否可以被另一个数整除
def divisible_by(num, den):
return (num % den) == 0
# 判断一个数是否为奇数
def is_odd(n):
return not divisible_by(n, 2)
# 获取时间步的嵌入向量
def get_timestep_embedding(timesteps, embedding_dim):
"""
此函数与 Denoising Diffusion Probabilistic Models 中的实现匹配：
来自 Fairseq。
构建正弦嵌入向量。
此实现与 tensor2tensor 中的实现匹配，但与 "Attention Is All You Need" 中第 3.5 节的描述略有不同。
"""
# 确保时间步的维度为 1
assert len(timesteps.shape) == 1
half_dim = embedding_dim // 2 # 计算嵌入维度的一半
emb = math.log(10000) / (half_dim - 1) # 计算对数因子
emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb) # 生成指数衰减的嵌入
emb = emb.to(device=timesteps.device) # 将嵌入移动到时间步的设备上
emb = timesteps.float()[:, None] * emb[None, :] # 扩展时间步并计算嵌入
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) # 计算正弦和余弦值
if embedding_dim % 2 == 1: # 如果嵌入维度为奇数，则进行零填充
emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
return emb # 返回嵌入
# 定义非线性激活函数，使用 Swish 激活函数
def nonlinearity(x):
# swish
return x * torch.sigmoid(x)
# 定义三维空间归一化模块
class SpatialNorm3D(nn.Module):
def __init__(
self,
f_channels, # 特征通道数
zq_channels, # 嵌入通道数
norm_layer=nn.GroupNorm, # 归一化层类型
freeze_norm_layer=False, # 是否冻结归一化层的参数
add_conv=False, # 是否添加卷积层
pad_mode="constant", # 填充模式
**norm_layer_params, # 归一化层的其他参数
):
super().__init__() # 调用父类构造函数
# 初始化归一化层
self.norm_layer = norm_layer(num_channels=f_channels, **norm_layer_params)
if freeze_norm_layer: # 如果需要冻结归一化层
for p in self.norm_layer.parameters: # 遍历所有参数
p.requires_grad = False # 不更新参数
self.add_conv = add_conv # 保存是否添加卷积层的标志
if self.add_conv: # 如果添加卷积层
# 创建三维因果卷积层
self.conv = CausalConv3d(zq_channels, zq_channels, kernel_size=3, pad_mode=pad_mode)
# 创建用于特征和嵌入的卷积层
self.conv_y = CausalConv3d(zq_channels, f_channels, kernel_size=1, pad_mode=pad_mode)
self.conv_b = CausalConv3d(zq_channels, f_channels, kernel_size=1, pad_mode=pad_mode)
# 前向传播函数
def forward(self, f, zq):
if zq.shape[2] > 1: # 如果嵌入的时间步数大于1
# 将特征拆分为第一时间步和剩余部分
f_first, f_rest = f[:, :, :1], f[:, :, 1:]
f_first_size, f_rest_size = f_first.shape[-3:], f_rest.shape[-3:] # 获取尺寸
zq_first, zq_rest = zq[:, :, :1], zq[:, :, 1:] # 拆分嵌入
# 使用最近邻插值调整 zq_first 的尺寸
zq_first = torch.nn.functional.interpolate(zq_first, size=f_first_size, mode="nearest")
# 使用最近邻插值调整 zq_rest 的尺寸
zq_rest = torch.nn.functional.interpolate(zq_rest, size=f_rest_size, mode="nearest")
zq = torch.cat([zq_first, zq_rest], dim=2) # 合并嵌入
else: # 如果时间步数为1
zq = torch.nn.functional.interpolate(zq, size=f.shape[-3:], mode="nearest") # 调整 zq 尺寸
if self.add_conv: # 如果添加卷积层
zq = self.conv(zq) # 通过卷积层处理 zq
norm_f = self.norm_layer(f) # 对特征进行归一化
# 计算新的特征值
new_f = norm_f * self.conv_y(zq) + self.conv_b(zq)
return new_f # 返回新的特征值
# 定义 Normalize3D 类的起始部分（未完成）
def Normalize3D(in_channels, zq_ch, add_conv):
# 返回一个三维空间归一化层的实例
return SpatialNorm3D(
# 输入通道数
in_channels,
# 量化通道数
zq_ch,
# 归一化层使用的类型，这里使用的是分组归一化
norm_layer=nn.GroupNorm,
# 是否冻结归一化层的参数，这里设置为不冻结
freeze_norm_layer=False,
# 是否添加卷积层，使用传入的参数
add_conv=add_conv,
# 归一化的组数
num_groups=32,
# 防止除零的极小值
eps=1e-6,
# 是否使用仿射变换，这里设置为使用
affine=True,
)
# 定义一个三维残差块类，继承自 nn.Module
class ResnetBlock3D(nn.Module):
# 初始化方法，接收多种参数以配置该块
def __init__(
self,
*,
in_channels, # 输入通道数
out_channels=None, # 输出通道数（可选）
conv_shortcut=False, # 是否使用卷积快捷连接
dropout, # dropout 比例
temb_channels=512, # 时间嵌入通道数
zq_ch=None, # zq 相关通道数（可选）
add_conv=False, # 是否添加卷积层
pad_mode="constant", # 填充模式
):
# 调用父类的初始化方法
super().__init__()
# 保存输入通道数
self.in_channels = in_channels
# 确定输出通道数，若未指定则等于输入通道数
out_channels = in_channels if out_channels is None else out_channels
# 保存输出通道数
self.out_channels = out_channels
# 保存是否使用卷积快捷连接的标志
self.use_conv_shortcut = conv_shortcut
# 初始化第一个归一化层
self.norm1 = Normalize3D(in_channels, zq_ch, add_conv=add_conv)
# 初始化第一个因果卷积层
self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3, pad_mode=pad_mode)
# 若时间嵌入通道数大于0，初始化时间嵌入投影层
if temb_channels > 0:
self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
# 初始化第二个归一化层
self.norm2 = Normalize3D(out_channels, zq_ch, add_conv=add_conv)
# 初始化 dropout 层
self.dropout = torch.nn.Dropout(dropout)
# 初始化第二个因果卷积层
self.conv2 = CausalConv3d(out_channels, out_channels, kernel_size=3, pad_mode=pad_mode)
# 如果输入和输出通道数不相等
if self.in_channels != self.out_channels:
# 如果使用卷积快捷连接，则初始化对应的卷积层
if self.use_conv_shortcut:
self.conv_shortcut = CausalConv3d(in_channels, out_channels, kernel_size=3, pad_mode=pad_mode)
# 否则，初始化 1x1 卷积层作为快捷连接
else:
self.nin_shortcut = torch.nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
# 前向传播方法
def forward(self, x, temb, zq):
# 将输入赋值给 h
h = x
# 对 h 进行归一化
h = self.norm1(h, zq)
# 应用非线性激活函数
h = nonlinearity(h)
# 通过第一个卷积层
h = self.conv1(h)
# 如果时间嵌入不为 None，则将其投影到 h 上
if temb is not None:
h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None, None]
# 对 h 进行第二次归一化
h = self.norm2(h, zq)
# 应用非线性激活函数
h = nonlinearity(h)
# 应用 dropout
h = self.dropout(h)
# 通过第二个卷积层
h = self.conv2(h)
# 如果输入和输出通道数不相等
if self.in_channels != self.out_channels:
# 如果使用卷积快捷连接
if self.use_conv_shortcut:
# 将输入 x 通过卷积快捷连接
x = self.conv_shortcut(x)
# 否则使用 1x1 卷积
else:
x = self.nin_shortcut(x)
# 返回输入和 h 的和
return x + h
# 定义一个二维注意力块类，继承自 nn.Module
class AttnBlock2D(nn.Module):
# 初始化方法，接收输入通道数、zq 通道数和是否添加卷积的标志
def __init__(self, in_channels, zq_ch=None, add_conv=False):
# 调用父类的初始化方法
super().__init__()
# 保存输入通道数
self.in_channels = in_channels
# 初始化归一化层
self.norm = Normalize3D(in_channels, zq_ch, add_conv=add_conv)
# 初始化查询、键、值卷积层，均为 1x1 卷积
self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
# 初始化输出卷积层
self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
# 前向传播函数，接受输入 x 和查询 zq
def forward(self, x, zq):
# 将输入 x 赋值给 h_
h_ = x
# 对 h_ 进行归一化处理，使用查询 zq
h_ = self.norm(h_, zq)
# 获取 h_ 的时间步长 t
t = h_.shape[2]
# 重排 h_ 的维度，将时间步和批次维度合并
h_ = rearrange(h_, "b c t h w -> (b t) c h w")
# 计算查询、键和值
q = self.q(h_) # 计算查询
k = self.k(h_) # 计算键
v = self.v(h_) # 计算值
# 计算注意力
b, c, h, w = q.shape # 解包 q 的形状信息
q = q.reshape(b, c, h * w) # 将 q 重塑为 (b, c, hw)
q = q.permute(0, 2, 1) # 变换维度顺序为 (b, hw, c)
k = k.reshape(b, c, h * w) # 将 k 重塑为 (b, c, hw)
# 计算 q 和 k 的批量矩阵乘法，得到注意力权重 w_
w_ = torch.bmm(q, k) # 计算注意力权重
w_ = w_ * (int(c) ** (-0.5)) # 对权重进行缩放
w_ = torch.nn.functional.softmax(w_, dim=2) # 对最后一维进行 softmax
# 根据注意力权重对值进行加权
v = v.reshape(b, c, h * w) # 将 v 重塑为 (b, c, hw)
w_ = w_.permute(0, 2, 1) # 变换维度顺序为 (b, hw, hw)
# 计算加权和，得到输出特征 h_
h_ = torch.bmm(v, w_) # 计算输出
h_ = h_.reshape(b, c, h, w) # 将 h_ 重塑回 (b, c, h, w)
# 对 h_ 进行投影
h_ = self.proj_out(h_)
# 将 h_ 的维度重排回原来的形状
h_ = rearrange(h_, "(b t) c h w -> b c t h w", t=t)
# 返回输入 x 和输出 h_ 的和
return x + h_
# 定义一个名为 MOVQDecoder3D 的神经网络模块类，继承自 nn.Module
class MOVQDecoder3D(nn.Module):
# 初始化方法，接受多个参数来配置解码器
def __init__(
self,
*,
ch, # 输入通道数
out_ch, # 输出通道数
ch_mult=(1, 2, 4, 8), # 通道数的倍增因子
num_res_blocks, # 残差块的数量
attn_resolutions, # 注意力机制的分辨率
dropout=0.0, # dropout 概率，用于防止过拟合
resamp_with_conv=True, # 是否使用卷积进行重采样
in_channels, # 输入的通道数
resolution, # 输入图像的分辨率
z_channels, # 噪声通道数
give_pre_end=False, # 是否给出预处理结束标志
zq_ch=None, # 量化通道数
add_conv=False, # 是否添加卷积层
pad_mode="first", # 填充模式，默认为 'first'
temporal_compress_times=4, # 时间压缩的倍数
**ignorekwargs, # 其他未指定的关键字参数
):
# 调用父类构造函数
super().__init__()
# 初始化通道数
self.ch = ch
# 初始化时间嵌入通道数为0
self.temb_ch = 0
# 获取分辨率数量
self.num_resolutions = len(ch_mult)
# 获取残差块数量
self.num_res_blocks = num_res_blocks
# 设置分辨率
self.resolution = resolution
# 设置输入通道数
self.in_channels = in_channels
# 设置是否在前向传播后给出结束标志
self.give_pre_end = give_pre_end
# 计算 temporal_compress_times 的 log2 值
self.temporal_compress_level = int(np.log2(temporal_compress_times))
# 如果 zq_ch 为 None，则使用 z_channels
if zq_ch is None:
zq_ch = z_channels
# 计算当前块的输入通道数
block_in = ch * ch_mult[self.num_resolutions - 1]
# 计算当前分辨率
curr_res = resolution // 2 ** (self.num_resolutions - 1)
# 设置 z 的形状
self.z_shape = (1, z_channels, curr_res, curr_res)
# 创建输入卷积层
self.conv_in = CausalConv3d(z_channels, block_in, kernel_size=3, pad_mode=pad_mode)
# 创建中间模块
self.mid = nn.Module()
# 添加第一个残差块
self.mid.block_1 = ResnetBlock3D(
in_channels=block_in,
out_channels=block_in,
temb_channels=self.temb_ch,
dropout=dropout,
zq_ch=zq_ch,
add_conv=add_conv,
pad_mode=pad_mode,
)
# 添加第二个残差块
self.mid.block_2 = ResnetBlock3D(
in_channels=block_in,
out_channels=block_in,
temb_channels=self.temb_ch,
dropout=dropout,
zq_ch=zq_ch,
add_conv=add_conv,
pad_mode=pad_mode,
)
# 创建上采样模块
self.up = nn.ModuleList()
# 从最高分辨率开始遍历
for i_level in reversed(range(self.num_resolutions)):
# 创建块和注意力模块的容器
block = nn.ModuleList()
attn = nn.ModuleList()
# 计算当前块的输出通道数
block_out = ch * ch_mult[i_level]
# 添加指定数量的残差块
for i_block in range(self.num_res_blocks + 1):
block.append(
ResnetBlock3D(
in_channels=block_in,
out_channels=block_out,
temb_channels=self.temb_ch,
dropout=dropout,
zq_ch=zq_ch,
add_conv=add_conv,
pad_mode=pad_mode,
)
)
# 更新输入通道数
block_in = block_out
# 如果当前分辨率在注意力分辨率列表中，添加注意力块
if curr_res in attn_resolutions:
attn.append(AttnBlock2D(block_in, zq_ch, add_conv=add_conv))
# 创建上采样模块
up = nn.Module()
up.block = block
up.attn = attn
# 如果不是最底层，进行上采样配置
if i_level != 0:
if i_level < self.num_resolutions - self.temporal_compress_level:
up.upsample = Upsample3D(block_in, resamp_with_conv, compress_time=False)
else:
up.upsample = Upsample3D(block_in, resamp_with_conv, compress_time=True)
# 更新当前分辨率为原来的两倍
curr_res = curr_res * 2
# 将上采样模块插入到列表开头以保持顺序
self.up.insert(0, up) # prepend to get consistent order
# 创建输出归一化层
self.norm_out = Normalize3D(block_in, zq_ch, add_conv=add_conv)
# 创建输出卷积层
self.conv_out = CausalConv3d(block_in, out_ch, kernel_size=3, pad_mode=pad_mode)
# 定义前向传播方法，接受输入 z 和可选参数 use_cp
def forward(self, z, use_cp=False):
# 保存输入 z 的形状以便后续使用
self.last_z_shape = z.shape
# 定义时间步嵌入变量，初始为 None
temb = None
# 获取 z 的时间步数（即第三维的大小）
t = z.shape[2]
# 将 z 赋值给 zq，用于后续计算
zq = z
# 通过输入卷积层处理 z
h = self.conv_in(z)
# 中间处理阶段
# 使用第一个中间块处理 h，传入 temb 和 zq
h = self.mid.block_1(h, temb, zq)
# h = self.mid.attn_1(h, zq) # 注释掉的注意力层
# 使用第二个中间块处理 h，传入 temb 和 zq
h = self.mid.block_2(h, temb, zq)
# 上采样阶段
# 反向遍历每个分辨率级别
for i_level in reversed(range(self.num_resolutions)):
# 遍历每个残差块
for i_block in range(self.num_res_blocks + 1):
# 在当前上采样级别中处理 h，传入 temb 和 zq
h = self.up[i_level].block[i_block](h, temb, zq)
# 如果当前级别有注意力层，则对 h 进行注意力处理
if len(self.up[i_level].attn) > 0:
h = self.up[i_level].attn[i_block](h, zq)
# 如果当前不是最后一个级别，则进行上采样
if i_level != 0:
h = self.up[i_level].upsample(h)
# 结束阶段
# 如果给定了预结束标志，则直接返回 h
if self.give_pre_end:
return h
# 对 h 进行归一化处理，传入 zq
h = self.norm_out(h, zq)
# 对 h 应用非线性激活函数
h = nonlinearity(h)
# 通过输出卷积层处理 h
h = self.conv_out(h)
# 返回最终的 h
return h
# 获取最后一层的卷积权重
def get_last_layer(self):
return self.conv_out.conv.weight
# 定义一个新的3D解码器类，继承自 nn.Module
class NewDecoder3D(nn.Module):
# 初始化方法，接受多个参数
def __init__(
self,
*,
ch, # 输入通道数
out_ch, # 输出通道数
ch_mult=(1, 2, 4, 8), # 通道数的倍增因子
num_res_blocks, # 残差块的数量
attn_resolutions, # 注意力分辨率
dropout=0.0, # dropout比率
resamp_with_conv=True, # 是否使用卷积进行重采样
in_channels, # 输入通道数
resolution, # 输入的分辨率
z_channels, # 噪声通道数
give_pre_end=False, # 是否返回预结束的输出
zq_ch=None, # 可选的量化通道数
add_conv=False, # 是否添加额外的卷积层
pad_mode="first", # 填充模式
temporal_compress_times=4, # 时间压缩次数
post_quant_conv=False, # 是否使用量化后的卷积
**ignorekwargs, # 其他忽略的参数
):
def forward(self, z):
# 断言输入的形状与预期的 z_shape 一致（已注释）
# self.last_z_shape = z.shape # 保存最后的输入形状
self.last_z_shape = z.shape
# 时间步嵌入初始化为 None
temb = None
# 获取输入 z 的时间步长
t = z.shape[2]
# 将 z 赋值给 zq 以备后续使用
zq = z
# 如果存在后量化卷积，则对 z 进行处理
if self.post_quant_conv is not None:
z = self.post_quant_conv(z)
# 对 z 进行初始卷积
h = self.conv_in(z)
# 中间处理阶段
h = self.mid.block_1(h, temb, zq) # 通过第一个中间块处理 h
# h = self.mid.attn_1(h, zq) # (注释掉)可能的注意力机制处理
h = self.mid.block_2(h, temb, zq) # 通过第二个中间块处理 h
# 上采样阶段
for i_level in reversed(range(self.num_resolutions)): # 反向遍历每个分辨率级别
for i_block in range(self.num_res_blocks + 1): # 遍历每个残差块
h = self.up[i_level].block[i_block](h, temb, zq) # 处理 h
# 如果有注意力模块，则应用注意力
if len(self.up[i_level].attn) > 0:
h = self.up[i_level].attn[i_block](h, zq)
# 如果当前级别不是0，则进行上采样
if i_level != 0:
h = self.up[i_level].upsample(h)
# 结束处理阶段
if self.give_pre_end: # 如果需要预结束输出，则返回 h
return h
# 通过归一化处理输出
h = self.norm_out(h, zq)
h = nonlinearity(h) # 应用非线性激活函数
h = self.conv_out(h) # 通过最终卷积层处理 h
return h # 返回最终的输出
# 获取最后一层的权重
def get_last_layer(self):
return self.conv_out.conv.weight # 返回卷积层的权重
.\cogvideo-finetune\sat\sgm\modules\autoencoding\vqvae\movq_dec_3d_dev.py
# pytorch_diffusion + derived encoder decoder
import math # 导入数学库，提供数学函数
import torch # 导入 PyTorch 库，进行张量计算
import torch.nn as nn # 导入 nn 模块，构建神经网络
import torch.nn.functional as F # 导入功能性模块，提供常用操作
import numpy as np # 导入 NumPy 库，进行数值计算
from beartype import beartype # 从 beartype 导入 beartype，用于类型检查
from beartype.typing import Union, Tuple, Optional, List # 导入类型提示
from einops import rearrange # 从 einops 导入 rearrange，用于重排张量维度
from .movq_enc_3d import CausalConv3d, Upsample3D, DownSample3D # 从本地模块导入 3D 卷积和上采样、下采样类
def cast_tuple(t, length=1):
# 如果 t 不是元组，则将其转换为指定长度的元组
return t if isinstance(t, tuple) else ((t,) * length)
def divisible_by(num, den):
# 检查 num 是否可以被 den 整除
return (num % den) == 0
def is_odd(n):
# 检查 n 是否为奇数
return not divisible_by(n, 2)
def get_timestep_embedding(timesteps, embedding_dim):
"""
这个函数构建正弦嵌入，与 Denoising Diffusion Probabilistic Models 的实现匹配：
来源于 Fairseq。
构建正弦嵌入。
与 tensor2tensor 的实现匹配，但与 "Attention Is All You Need" 第 3.5 节中的描述略有不同。
"""
# 确保 timesteps 是一维的
assert len(timesteps.shape) == 1
# 计算嵌入维度的一半
half_dim = embedding_dim // 2
# 计算嵌入的基础
emb = math.log(10000) / (half_dim - 1)
# 计算正弦嵌入的值
emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
# 将嵌入移动到 timesteps 的设备上
emb = emb.to(device=timesteps.device)
# 计算每个时间步的嵌入
emb = timesteps.float()[:, None] * emb[None, :]
# 将正弦和余弦值连接起来
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
# 如果嵌入维度是奇数，进行零填充
if embedding_dim % 2 == 1: # zero pad
emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
# 返回最终的嵌入
return emb
def nonlinearity(x):
# 使用 Swish 激活函数
return x * torch.sigmoid(x)
class SpatialNorm3D(nn.Module):
# 定义一个 3D 空间归一化的类
def __init__(
self,
f_channels,
zq_channels,
norm_layer=nn.GroupNorm,
freeze_norm_layer=False,
add_conv=False,
pad_mode="constant",
**norm_layer_params,
):
# 初始化函数，设置参数
super().__init__() # 调用父类的初始化函数
# 创建归一化层
self.norm_layer = norm_layer(num_channels=f_channels, **norm_layer_params)
# 如果需要冻结归一化层的参数
if freeze_norm_layer:
for p in self.norm_layer.parameters: # 遍历归一化层的参数
p.requires_grad = False # 冻结参数不进行更新
self.add_conv = add_conv # 是否添加卷积层
# 如果添加卷积层，创建 causal 卷积层
if self.add_conv:
# self.conv = nn.Conv3d(zq_channels, zq_channels, kernel_size=3, stride=1, padding=1)
self.conv = CausalConv3d(zq_channels, zq_channels, kernel_size=3, pad_mode=pad_mode)
# 创建一个 1x1 卷积层用于 y 和 b 通道
# self.conv_y = nn.Conv3d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0)
# self.conv_b = nn.Conv3d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0)
self.conv_y = CausalConv3d(zq_channels, f_channels, kernel_size=1, pad_mode=pad_mode)
self.conv_b = CausalConv3d(zq_channels, f_channels, kernel_size=1, pad_mode=pad_mode)
# 定义前向传播方法，接收输入 f 和 zq
def forward(self, f, zq):
# 如果 zq 的第三维大于 1，表示有多个通道
if zq.shape[2] > 1:
# 分割 f 为第一个通道和其余通道
f_first, f_rest = f[:, :, :1], f[:, :, 1:]
# 获取第一个通道和其余通道的尺寸
f_first_size, f_rest_size = f_first.shape[-3:], f_rest.shape[-3:]
# 分割 zq 为第一个通道和其余通道
zq_first, zq_rest = zq[:, :, :1], zq[:, :, 1:]
# 对第一个通道进行最近邻插值调整尺寸
zq_first = torch.nn.functional.interpolate(zq_first, size=f_first_size, mode="nearest")
# 对其余通道进行最近邻插值调整尺寸
zq_rest = torch.nn.functional.interpolate(zq_rest, size=f_rest_size, mode="nearest")
# 将调整后的通道合并在一起
zq = torch.cat([zq_first, zq_rest], dim=2)
# 如果 zq 只有一个通道，直接调整其尺寸
else:
zq = torch.nn.functional.interpolate(zq, size=f.shape[-3:], mode="nearest")
# 如果需要添加卷积层
if self.add_conv:
# 对 zq 进行卷积操作
zq = self.conv(zq)
# 对 f 应用归一化层
norm_f = self.norm_layer(f)
# 计算新的 f，结合卷积后的 zq
new_f = norm_f * self.conv_y(zq) + self.conv_b(zq)
# 返回新的 f
return new_f
# 定义一个 3D 归一化函数，接收输入通道、量化通道及是否添加卷积的标志
def Normalize3D(in_channels, zq_ch, add_conv):
# 调用空间归一化 3D，传入相应参数
return SpatialNorm3D(
in_channels,
zq_ch,
norm_layer=nn.GroupNorm, # 使用分组归一化层
freeze_norm_layer=False, # 不冻结归一化层
add_conv=add_conv, # 是否添加卷积
num_groups=32, # 设置组的数量
eps=1e-6, # 设置小常数以防止除零
affine=True, # 使用仿射变换
)
# 定义 3D ResNet 块，继承自 nn.Module
class ResnetBlock3D(nn.Module):
# 初始化方法，接收多个参数配置
def __init__(
self,
*,
in_channels, # 输入通道数
out_channels=None, # 输出通道数，可选
conv_shortcut=False, # 是否使用卷积短接
dropout, # dropout 比率
temb_channels=512, # 时间嵌入通道数
zq_ch=None, # 量化通道
add_conv=False, # 是否添加卷积
pad_mode="constant", # 填充模式
):
super().__init__() # 调用父类构造函数
self.in_channels = in_channels # 设置输入通道数
out_channels = in_channels if out_channels is None else out_channels # 设置输出通道数
self.out_channels = out_channels # 保存输出通道数
self.use_conv_shortcut = conv_shortcut # 保存是否使用卷积短接的标志
# 创建第一个归一化层
self.norm1 = Normalize3D(in_channels, zq_ch, add_conv=add_conv)
# self.conv1 = torch.nn.Conv3d(in_channels,
# out_channels,
# kernel_size=3,
# stride=1,
# padding=1)
# 使用因果卷积创建第一个卷积层
self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3, pad_mode=pad_mode)
if temb_channels > 0: # 如果时间嵌入通道数大于零
# 创建时间嵌入投影层
self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
# 创建第二个归一化层
self.norm2 = Normalize3D(out_channels, zq_ch, add_conv=add_conv)
# 创建 dropout 层
self.dropout = torch.nn.Dropout(dropout)
# self.conv2 = torch.nn.Conv3d(out_channels,
# out_channels,
# kernel_size=3,
# stride=1,
# padding=1)
# 使用因果卷积创建第二个卷积层
self.conv2 = CausalConv3d(out_channels, out_channels, kernel_size=3, pad_mode=pad_mode)
# 如果输入和输出通道数不一致
if self.in_channels != self.out_channels:
if self.use_conv_shortcut: # 如果使用卷积短接
# self.conv_shortcut = torch.nn.Conv3d(in_channels,
# out_channels,
# kernel_size=3,
# stride=1,
# padding=1)
# 使用因果卷积创建短接卷积层
self.conv_shortcut = CausalConv3d(in_channels, out_channels, kernel_size=3, pad_mode=pad_mode)
else:
# 创建一个 1x1 的卷积层作为短接
self.nin_shortcut = torch.nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
# self.nin_shortcut = CausalConv3d(in_channels, out_channels, kernel_size=1, pad_mode=pad_mode)
# 定义前向传播函数，接收输入张量 x、时间嵌入 temb 和 zq
def forward(self, x, temb, zq):
# 将输入赋值给 h
h = x
# 对 h 应用第一层归一化，使用 zq 作为参数
h = self.norm1(h, zq)
# 对 h 应用非线性激活函数
h = nonlinearity(h)
# 对 h 应用第一层卷积
h = self.conv1(h)
# 如果时间嵌入 temb 不为空
if temb is not None:
# 将时间嵌入经过非线性处理并进行投影后加到 h
h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None, None]
# 对 h 应用第二层归一化，使用 zq 作为参数
h = self.norm2(h, zq)
# 对 h 应用非线性激活函数
h = nonlinearity(h)
# 对 h 应用 dropout 操作
h = self.dropout(h)
# 对 h 应用第二层卷积
h = self.conv2(h)
# 如果输入通道数与输出通道数不相等
if self.in_channels != self.out_channels:
# 如果使用卷积短路
if self.use_conv_shortcut:
# 对输入 x 应用卷积短路
x = self.conv_shortcut(x)
else:
# 对输入 x 应用 NIN 短路
x = self.nin_shortcut(x)
# 返回输入 x 与 h 的和
return x + h
# 定义一个二维注意力块类，继承自 nn.Module
class AttnBlock2D(nn.Module):
# 初始化方法，设置输入通道数、zq 通道数以及是否添加卷积
def __init__(self, in_channels, zq_ch=None, add_conv=False):
# 调用父类的初始化方法
super().__init__()
# 保存输入通道数
self.in_channels = in_channels
# 初始化 3D 归一化层
self.norm = Normalize3D(in_channels, zq_ch, add_conv=add_conv)
# 定义查询（Q）卷积层，kernel_size=1
self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
# 定义键（K）卷积层，kernel_size=1
self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
# 定义值（V）卷积层，kernel_size=1
self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
# 定义输出投影卷积层，kernel_size=1
self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
# 前向传播方法
def forward(self, x, zq):
# 保存输入数据
h_ = x
# 对输入数据进行归一化
h_ = self.norm(h_, zq)
# 获取时间步长
t = h_.shape[2]
# 重新排列张量维度，将时间步和批次合并
h_ = rearrange(h_, "b c t h w -> (b t) c h w")
# 计算查询（Q）、键（K）和值（V）
q = self.q(h_)
k = self.k(h_)
v = self.v(h_)
# 计算注意力
b, c, h, w = q.shape # 获取批次大小、通道数、高度和宽度
q = q.reshape(b, c, h * w) # 重新排列查询张量维度
q = q.permute(0, 2, 1) # 交换维度顺序，变为 b, hw, c
k = k.reshape(b, c, h * w) # 重新排列键张量维度
# 计算注意力权重矩阵，使用批量矩阵乘法
w_ = torch.bmm(q, k) # b, hw, hw，计算 q 和 k 的点积
w_ = w_ * (int(c) ** (-0.5)) # 对权重进行缩放
w_ = torch.nn.functional.softmax(w_, dim=2) # 对权重进行 softmax 归一化
# 注意力机制应用于值（V）
v = v.reshape(b, c, h * w) # 重新排列值张量维度
w_ = w_.permute(0, 2, 1) # 交换权重维度顺序
# 计算加权值，得到注意力输出
h_ = torch.bmm(v, w_) # b, c, hw，计算 v 和权重的点积
h_ = h_.reshape(b, c, h, w) # 重新排列输出张量维度
# 通过输出投影层进行变换
h_ = self.proj_out(h_)
# 恢复张量到原来的维度结构
h_ = rearrange(h_, "(b t) c h w -> b c t h w", t=t)
# 返回输入和输出的和
return x + h_
# 定义一个三维 MOVQ 解码器类，继承自 nn.Module
class MOVQDecoder3D(nn.Module):
# 初始化方法，设置多个参数，包括通道数、分辨率等
def __init__(
self,
*,
ch,
out_ch,
ch_mult=(1, 2, 4, 8),
num_res_blocks,
attn_resolutions,
dropout=0.0,
resamp_with_conv=True,
in_channels,
resolution,
z_channels,
give_pre_end=False,
zq_ch=None,
add_conv=False,
pad_mode="first",
temporal_compress_times=4,
**ignorekwargs,
# 定义前向传播函数，接受输入 z 和一个可选的 use_cp 参数
def forward(self, z, use_cp=False):
# 断言输入 z 的形状与预期的形状一致（此行被注释掉）
# assert z.shape[1:] == self.z_shape[1:]
# 保存输入 z 的形状，以备后用
self.last_z_shape = z.shape
# 初始化时间步嵌入变量
temb = None
# 获取 z 的时间步长度
t = z.shape[2]
# 将 z 赋值给 zq，作为输入块
zq = z
# 对输入 z 进行卷积操作，得到初步特征 h
h = self.conv_in(z)
# 中间层处理
h = self.mid.block_1(h, temb, zq) # 通过第一个中间块处理特征
# h = self.mid.attn_1(h, zq) # 注释掉的注意力机制处理
h = self.mid.block_2(h, temb, zq) # 通过第二个中间块处理特征
# 上采样过程
for i_level in reversed(range(self.num_resolutions)): # 反向遍历分辨率层级
for i_block in range(self.num_res_blocks + 1): # 遍历每个块
h = self.up[i_level].block[i_block](h, temb, zq) # 对当前特征进行块处理
if len(self.up[i_level].attn) > 0: # 如果当前层有注意力机制
h = self.up[i_level].attn[i_block](h, zq) # 进行注意力机制处理
if i_level != 0: # 如果不是最后一层
h = self.up[i_level].upsample(h) # 进行上采样
# 结束处理
if self.give_pre_end: # 如果需要返回中间结果
return h
h = self.norm_out(h, zq) # 对 h 进行规范化处理
h = nonlinearity(h) # 应用非线性激活函数
h = self.conv_out(h) # 最后卷积操作，输出结果
return h # 返回最终结果
# 获取最后一层的权重
def get_last_layer(self):
return self.conv_out.conv.weight # 返回卷积层的权重
# 定义一个新的 3D 解码器类，继承自 nn.Module
class NewDecoder3D(nn.Module):
# 初始化方法，接收多个参数以配置解码器
def __init__(
self,
*,
ch, # 输入通道数
out_ch, # 输出通道数
ch_mult=(1, 2, 4, 8), # 通道倍增因子
num_res_blocks, # 残差块数量
attn_resolutions, # 注意力分辨率
dropout=0.0, # dropout 比率
resamp_with_conv=True, # 是否使用卷积进行重采样
in_channels, # 输入通道
resolution, # 输入分辨率
z_channels, # 噪声通道数
give_pre_end=False, # 是否给出预处理结束
zq_ch=None, # 可选的量化通道数
add_conv=False, # 是否添加额外卷积层
pad_mode="first", # 填充模式
temporal_compress_times=4, # 时间压缩倍数
post_quant_conv=False, # 是否使用后量化卷积
**ignorekwargs, # 其他忽略的关键字参数
):
# 初始化父类 nn.Module
super(NewDecoder3D, self).__init__()
# 定义前向传播方法，接收输入 z
def forward(self, z):
# 断言检查 z 的形状是否与 z_shape 匹配
# assert z.shape[1:] == self.z_shape[1:]
# 记录 z 的最后形状
self.last_z_shape = z.shape
# 定义时间步嵌入，初始为 None
temb = None
# 获取 z 的时间步长
t = z.shape[2]
# z 赋值给 zq，作为量化输入
zq = z
# 如果定义了后量化卷积，则对 z 进行处理
if self.post_quant_conv is not None:
z = self.post_quant_conv(z)
# 对输入 z 进行初始卷积处理
h = self.conv_in(z)
# 中间层处理
h = self.mid.block_1(h, temb, zq) # 通过第一个中间块
# h = self.mid.attn_1(h, zq) # 可选的注意力机制
h = self.mid.block_2(h, temb, zq) # 通过第二个中间块
# 上采样处理
for i_level in reversed(range(self.num_resolutions)): # 从高到低分辨率处理
for i_block in range(self.num_res_blocks + 1): # 遍历每个残差块
h = self.up[i_level].block[i_block](h, temb, zq) # 通过当前上采样块
# 如果存在注意力模块，则应用
if len(self.up[i_level].attn) > 0:
h = self.up[i_level].attn[i_block](h, zq)
# 如果不是最后一层，则进行上采样
if i_level != 0:
h = self.up[i_level].upsample(h)
# 结束处理
if self.give_pre_end: # 如果需要预处理结束，则返回 h
return h
# 对输出进行归一化处理
h = self.norm_out(h, zq)
# 应用非线性激活函数
h = nonlinearity(h)
# 通过最终卷积层得到输出
h = self.conv_out(h)
# 返回最终输出
return h
# 定义获取最后一层权重的方法
def get_last_layer(self):
# 返回最后卷积层的权重
return self.conv_out.conv.weight
