当前位置：首页 > article >正文

AF3 FourierEmbedding类源码解读

article 2025/1/27 17:08:25

FourierEmbedding 是一个用于扩散条件的傅里叶嵌入类，其核心是将输入的时间步噪声强度或控制参数（timestep）转换为高维的周期性特征。

源代码：

class FourierEmbedding(nn.Module):
    """Fourier embedding for diffusion conditioning."""
    def __init__(self, embed_dim):
        super(FourierEmbedding, self).__init__()
        self.embed_dim = embed_dim
        # Randomly generate weight/bias once before training
        self.weight = nn.Parameter(torch.randn((1, embed_dim)))
        self.bias = nn.Parameter(torch.randn((1, embed_dim)))

    def forward(self, t):
        """Compute embeddings"""
        two_pi = torch.tensor(2 * 3.1415, device=t.device, dtype=t.dtype)
        return torch.cos(two_pi * (t * self.weight + self.bias))

类代码解读：

1. 类的功能

该模块的主要目的是通过傅里叶变换，将输入的时间步嵌入到一个周期性的高维特征空间。这种处理方式在扩散模型中尤为重要，因为时间步本身是一个标量（单一数值），通过傅里叶嵌入，模型能够更好地捕获时间的周期性模式。

2. `init` 方法

def __init__(self, embed_dim):
    super(FourierEmbedding, self).__init__()
    self.embed_dim = embed_dim
    # Randomly generate weight/bias once before training
    self.weight = nn.Parameter(torch.randn((1, embed_dim)))
    self.bias = nn.Parameter(torch.randn((1, embed_dim)))