mindspore打卡第六天DDPM 中的attention部分

mindspore打卡第六天DDPM 中的attention部分

本文我们同样假设神经网络只需要学习（表示）这个条件概率分布的平均值 𝜇𝜃 目标函数来学习反向过程的平均值，作者观察到 𝑞和 𝑝𝜃的组合可以被视为变分自动编码器(VAE)

为了导出一个目标函数来学习反向过程的平均值，作者观察到 $q$ 和 $p_\theta$ 的组合可以被视为变分自动编码器(VAE)。因此，变分下界（也称为ELBO）可用于最小化真值数据样本 $\mathbf{x}_0$ 的似然负对数（有关ELBO的详细信息，请参阅VAE论文[（Kingma等人，2013年）](https://arxiv.org/abs/1312.6114)），该过程的ELBO是每个时间步长的损失之和 $L=L_0+L_1+...+L_T$ ，其中，每项的损失 $L_t$ （除了 $L_0$ ）实际上是2个高斯分布之间的KL发散，可以明确地写为相对于均值的L2-loss!

如Sohl-Dickstein等人所示，构建Diffusion正向过程的直接结果是我们可以在条件是 $\mathbf{x}_0$ （因为高斯和也是高斯）的情况下，在任意噪声水平上采样 $\mathbf{x}_t$ ，而不需要重复应用 $q$ 去采样 $\mathbf{x}_t$ ，这非常方便。使用

$$
\\\alpha_t := 1 - \beta_t\\\\\bar{\alpha}t := \Pi_{s=1}^{t} \alpha_s\\
$$

我们就有

$$  
q(\mathbf{x}_t | \mathbf{x}_0) = \cal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1- \bar{\alpha}_t) \mathbf{I})
$$

这意味着我们可以采样高斯噪声并适当地缩放它，然后将其添加到 $\mathbf{x}_0$ 中，直接获得 $\mathbf{x}_t$ 。

请注意，$\bar{\alpha}_t$ 是已知 $\beta_t$ 方差计划的函数，因此也是已知的，可以预先计算。这允许我们在训练期间优化损失函数 $L$ 的随机项。或者换句话说，在训练期间随机采样 $t$ 并优化 $L_t$ 。

##  $\alpha_t$ 是X0信号的保留程度

正如Ho等人所展示的那样，这种性质的另一个优点是可以重新参数化平均值，使神经网络学习（预测）构成损失的KL项中噪声的附加噪声。这意味着我们的神经网络变成了噪声预测器，而不是（直接）均值预测器。其中，平均值可以按如下方式计算：

$$ \mathbf{\mu}_\theta(\mathbf{x}_t, t) = \frac{1}{\sqrt{\alpha_t}} \left(  \mathbf{x}_t - \frac{\beta_t}{\sqrt{1- \bar{\alpha}_t}} \mathbf{\epsilon}_\theta(\mathbf{x}_t, t) \right) $$

最终的目标函数 ${L}_{t}$ 如下 （随机步长 t 由 $({\epsilon} \sim N(\mathbf{0}, \mathbf{I}))$ 给定）：

$$ \| \mathbf{\epsilon} - \mathbf{\epsilon}_\theta(\mathbf{x}_t, t) \|^2 = \| \mathbf{\epsilon} - \mathbf{\epsilon}_\theta( \sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{(1- \bar{\alpha}_t)  } \mathbf{\epsilon}, t) \|^2$$

在训练这样的模型时，实际添加的噪声 $\mathbf{\epsilon}$ 是通过数学构造或模拟过程得到的，而不是从实际数据中直接测量。具体来说，在扩散模型的训练阶段，这一过程通常涉及以下几个步骤：

1. **数据准备**：首先，你需要原始数据集，比如一组图像或其他类型的数据。

2. **扩散过程的定义**：然后定义一个前向扩散过程，这个过程逐步向原始数据中添加噪声。对于每个时间步长 $t$，会有关联的噪声水平参数（如 $\beta_t$），这些参数控制着在该时间步添加到数据中的噪声量。随着 $t$ 增加，数据点 $\mathbf{x}_0$ 会变得越来越模糊，直到最终接近纯粹的噪声。

3. **噪声的合成**：对于训练中的每一个样本 $\mathbf{x}_0$ 和选定的时间步 $t$，实际的噪声 $\mathbf{\epsilon}$ 是从一个预先定义的噪声分布中抽样得到的，通常是高斯分布 $N(\mathbf{0}, \mathbf{I})$。这意味着对于每次训练迭代，都会为特定的时间步和样本生成一个新的噪声向量，这个噪声向量与原始数据结合，通过一定的公式（比如上面提到的方程）向前扩散，得到带有噪声的数据表示 $\mathbf{x}_t$。

4. **目标函数的构建**：训练的目标是让神经网络 $\mathbf{\epsilon}_\theta(\mathbf{x}_t, t)$ 学习预测这个已知的、添加的噪声 $\mathbf{\epsilon}$。这是通过比较网络预测的噪声与真实添加的噪声，并优化一个如均方误差（MSE）之类的损失函数来实现的。

简而言之，$\mathbf{\epsilon}$ 在训练中是通过从标准正态分布中随机抽样获得的，作为理论上的“真实”噪声添加到数据中，以构建训练实例。模型随后学习逆向此过程，即从含噪数据中预测或逆推原本的噪声分量，从而能够学习到数据的生成过程。

U-Net确实可以从高层次上被视为一系列卷积层（用于下采样）和反卷积层（用于上采样）的组合，中间穿插有跳跃连接（skip connections），这些连接允许来自编码器部分的特征图直接传递到相应层级的解码器部分。这种设计不仅简化了特征从低层到高层再回来的路径，而且有助于保留对细节的精确定位，这对于许多像素级预测任务至关重要，例如图像分割或噪声预测。

**Bottleneck Layer的理解：**

在U-Net中，虽然没有明确标记为“bottleneck”的单一层，但可以认为网络的“狭窄”部分——即从最大下采样层到开始上采样的转折点——扮演了类似自动编码器中bottleneck的角色。在这个区域，网络的表征最为紧凑，因为它在编码器部分经过了多次下采样操作，将输入图像的信息浓缩到了较小的空间维度上。这个“狭窄”部分强制网络学习数据的关键特征，而忽略不重要的细节，类似于自动编码器中的信息瓶颈，其目的是提取数据的核心特征表示。

**残差连接的作用：**

虽然原始U-Net论文中并未直接提及残差学习（Residual Learning），后者是Deep Residual Network (ResNet) 提出的概念，但U-Net的设计中确实通过跳跃连接实现了类似的功能。这些跳跃连接使得网络能够直接将早期层的特征与后期层的特征相结合，帮助缓解了梯度消失问题，并促进了深层网络的优化。这与残差学习的目的——通过学习输入到输出的残差（而非直接学习输出本身）来简化网络训练——有异曲同工之妙，尽管两者在具体实现细节上有所区别。

综上所述，U-Net不仅仅是一个简单的卷积层或反卷积层的堆叠，它的独特之处在于加入了跳跃连接，这些连接允许网络高效地利用多尺度信息，这对于精确的图像处理任务特别有效，包括噪声预测任务。

## 构建Diffusion模型

下面，我们逐步构建Diffusion模型。

首先，我们定义了一些帮助函数和类，这些函数和类将在实现神经网络时使用。

这段代码定义了一系列辅助函数和一个类，主要用于深度学习模型中的一些常见操作，特别是与注意力机制和残差连接相关的处理。下面是对每个函数和类的中文注释说明：

### 函数 `rearrange`
```python
def rearrange(head, inputs):
    b, hc, x, y = inputs.shape
    c = hc // head
    return inputs.reshape((b, head, c, x * y))
```
**功能**：根据`head`数重排张量形状，常用于注意力机制中对多头注意力的处理。它将输入张量按照头数重新排列，使得张量形状适配多头注意力的计算需求。

### 函数 `rsqrt`
```python
def rsqrt(x):
    res = ops.sqrt(x)
    return ops.inv(res)
```
**功能**：计算输入张量元素的平方根的倒数。这在某些算法中用于归一化或初始化权重时可能会用到。

### 函数 `randn_like`
```python
def randn_like(x, dtype=None):
    # ...省略的代码逻辑与randn相同
```
**功能**：生成与给定张量`x`形状相同且数据类型可选的随机正态分布（标准差为1，均值为0）张量。

### 函数 `randn`
```python
def randn(shape, dtype=None):
    # ...省略的代码逻辑与randn_like中的实现相同，只是直接根据shape生成
```
**功能**：根据给定的形状和数据类型（默认为float32）生成随机正态分布张量。

### 函数 `randint`
```python
def randint(low, high, size, dtype=ms.int32):
    # ...通过ops.uniform生成指定范围内的随机整数张量
```
**功能**：生成指定区间内的随机整数张量，区间为`[low, high)`，大小和数据类型均可指定。

### 函数 `exists`
```python
def exists(x):
    return x is not None
```
**功能**：检查变量是否非空，返回True或False。

### 函数 `default`
```python
def default(val, d):
    # ...如果val存在则返回val，否则根据d的类型返回默认值或调用默认值的函数
```
**功能**：提供一个值或默认值，如果提供的值存在则直接返回，否则根据第二个参数d的类型返回默认值或调用默认值对应的函数。

### 函数 `_check_dtype`
```python
def _check_dtype(d1, d2):
    # ...判断两个数据类型是否兼容，优先返回float32或共同的数据类型，否则抛出异常
```
**功能**：检查两个数据类型是否兼容，确保运算中数据类型的正确性，避免不支持的数据类型组合。

### 类 `Residual`
```python
class Residual(nn.Cell):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def construct(self, x, *args, **kwargs):
        return self.fn(x, *args, **kwargs) + x
```
**功能**：实现残差连接的神经网络层。这个类包装了一个函数`fn`，在调用时会将函数的输出与输入`x`相加，形成残差连接，常用于深度学习模型中增加网络的训练稳定性及深度。

我们还定义了上采样和下采样操作的别名

def Upsample(dim):
    return nn.Conv2dTranspose(dim, dim, 4, 2, pad_mode="pad", padding=1)

def Downsample(dim):
    return nn.Conv2d(dim, dim, 4, 2, pad_mode="pad", padding=1)
    
卷积层和下采样操作在深度学习中都用于处理输入数据的空间维度，但它们的目的和方法有所不同。

**卷积层（包括上采样和下采样卷积）**：
- **上采样（Upsampling）**：通过卷积 transpose（转置卷积），也称为反卷积，可以实现上采样。上采样操作增加了特征图的空间尺寸，通常用于生成任务中重建较高分辨率的图像或作为部分网络架构中的解码器部分，以从较低层次的特征中恢复细节。例如，`nn.Conv2dTranspose(dim, dim, 4, 2, pad_mode="pad", padding=1)` 这行代码定义了一个转置卷积层，它使用4x4的卷积核，步长为2，进行上采样，使得输出的尺寸是输入尺寸的两倍，并在边缘使用了1像素的padding以保持输出尺寸的对称性。

- **下采样（Downsampling）**：虽然传统上我们可能认为下采样主要通过最大池化（MaxPooling）或平均池化（AveragePooling）等操作完成，但实际上，卷积层也可以用来执行空间下采样。通过设置合适的步长（stride）大于1，卷积层能够在提取特征的同时减少空间维度。在给出的代码`nn.Conv2d(dim, dim, 4, 2, pad_mode="pad", padding=1)`中，尽管使用的是一般的卷积操作（`nn.Conv2d`），但由于设置了步长为2，实际上它也起到了下采样的作用，减小了输入特征图的空间尺寸，同时保持了通道数不变。这种操作常用于编码器部分，帮助减少计算负担并增加模型的泛化能力。

**池化层与下采样的关系**：
- 池化层（如最大池化或平均池化）是专门设计用于下采样的操作，其主要目的是降低空间维度，减少计算复杂度，并增加一定程度的平移不变性。池化不改变通道数，但会显著减少特征图的空间分辨率，通常不涉及学习参数。

总结来说，虽然下采样通常与池化层相关联，卷积层（特别是当设置步长大于1时）也能实现下采样功能。而上采样则主要通过转置卷积来完成，目的是放大特征图尺寸，为生成高分辨率输出或在模型中构建上采样路径提供支持。

### 位置向量

由于神经网络的参数在时间（噪声水平）上共享，作者使用正弦位置嵌入来编码$t$，灵感来自Transformer（[Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)）。对于批处理中的每一张图像，神经网络"知道"它在哪个特定时间步长（噪声水平）上运行。

`SinusoidalPositionEmbeddings`模块采用`(batch_size, 1)`形状的张量作为输入（即批处理中几个有噪声图像的噪声水平），并将其转换为`(batch_size, dim)`形状的张量，其中`dim`是位置嵌入的尺寸。然后，我们将其添加到每个剩余块中。

class SinusoidalPositionEmbeddings(nn.Cell):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = np.exp(np.arange(half_dim) * - emb)
        self.emb = Tensor(emb, ms.float32)

    def construct(self, x):
        emb = x[:, None] * self.emb[None, :]
        emb = ops.concat((ops.sin(emb), ops.cos(emb)), axis=-1)  ###前面5个是sin 后面5个是cos
        return emb

# 初始化实例
pos_embedding = SinusoidalPositionEmbeddings(dim=10)

# 假设我们想为前4个位置生成位置嵌入
positions = ms.Tensor(np.array([0, 1, 2, 3]), ms.int32)

# 通过构造函数得到位置嵌入
output_embeddings = pos_embedding(positions)

# 打印输入和输出数据
print("Input positions:", positions)
print("Output embeddings shape:", output_embeddings.shape)

在 `SinusoidalPositionEmbeddings` 类中实现的位置编码方法源于 Vaswani 等人在 “Attention is All You Need” 论文中提出的方法。其核心思想是为输入序列中的每个位置 \( p \) 生成一个固定的向量表示，该向量由正弦和余弦函数的值构成，以此捕获位置信息。对于维度 \( d \)（其中 \( d \) 为偶数，以保证可以平分为正弦和余弦部分），位置 \( p \) 的嵌入向量可以通过以下公式获得：

$$\[ PE(p, 2k) = \sin\left(\frac{p}{10000^{2k/D}}\right) \]$$
$$\[ PE(p, 2k+1) = \cos\left(\frac{p}{10000^{2k/D}}\right) \]$$

其中：
- \( PE(p, 2k) \) 和 \( PE(p, 2k+1) \) 分别代表位置 \( p \) 在维度 \( 2k \) 和 \( 2k+1 \) 上的正弦和余弦位置嵌入值。
- \( D \) 是位置嵌入的总维度，即类初始化时的 `dim` 参数。
- \( k \) 是维度的一半之前（因此是0到\( \frac{D}{2}-1 \)的整数）的索引。
- \( p \) 是位置索引（从0开始）。

在初始化阶段，类计算了一个基于上述公式的频率标量数组 `emb`，这些标量随后与位置索引相乘，以生成特定于位置的正弦和余弦值。在 `construct` 方法中，通过计算正弦和余弦值，然后沿最后一个轴拼接它们，从而为每个位置生成完整的嵌入向量。

因此，对于维度 \( dim=10 \)，前5个维度是基于位置的正弦函数值，后5个维度是基于同一位置的余弦函数值。每个维度上的计算依据上述公式，确保了不同位置在嵌入空间中有独特的表示，且这种表示方式能够编码相对位置信息，对模型理解序列数据非常有用。

np.exp(np.arange(4) * - 5)

在Python中使用NumPy库，`np.exp()`函数用于计算自然指数\(e\)的幂。`np.arange(4)`生成一个包含从0到3的整数数组，而`* -5`则是将数组中的每个元素乘以-5。

下面是`np.exp(np.arange(4) * -5)`的具体运算步骤：

1. `np.arange(4)`生成一个数组：`[0, 1, 2, 3]`。
2. 将这个数组中的每个元素乘以-5：`[0 * -5, 1 * -5, 2 * -5, 3 * -5]`，得到数组`[-0, -5, -10, -15]`。
3. 使用`np.exp()`对上一步得到的数组中的每个元素计算\(e\)的指数幂。即计算\(e\)的0次幂、\(e\)的-5次幂、\(e\)的-10次幂和\(e\)的-15次幂。

具体计算结果如下：

- \( e^0 = 1 \)
- \( e^{-5} \approx 0.0067 \)（这是一个近似值）
- \( e^{-10} \approx 0.0001 \)（这是一个近似值）
- \( e^{-15} \approx 3.2e-7 \)（这是一个近似值）

最终，`np.exp(np.arange(4) * -5)`的结果将是一个包含上述计算结果的数组。这个数组的每个元素都是对应指数幂的值。

ConvNeXT（Liu et al., 2022）替换ResNet

class Block(nn.Cell):
    def __init__(self, dim, dim_out, groups=1):
        super().__init__()
        self.proj = nn.Conv2d(dim, dim_out, 3, pad_mode="pad", padding=1)
        self.proj =         c(dim, dim_out, 3, padding=1, pad_mode='pad')  ####在类内强调用c执行者这个步骤？  本质是一个卷积层
        self.norm = nn.GroupNorm(groups, dim_out)
        self.act = nn.SiLU()

    def construct(self, x, scale_shift=None):
        x = self.proj(x)
        x = self.norm(x)

        if exists(scale_shift):
            scale, shift = scale_shift
            x = x * (scale + 1) + shift

        x = self.act(x)
        return x

class ConvNextBlock(nn.Cell):
    def __init__(self, dim, dim_out, *, time_emb_dim=None, mult=2, norm=True):
        super().__init__()
        self.mlp = (
            nn.SequentialCell(nn.GELU(), nn.Dense(time_emb_dim, dim))
            if exists(time_emb_dim)  ##检查变量非空  如果 time_emb_dim 未提供，则 mlp 为 None 否则创建一个多层感知机（MLP）
            else None
        )

        self.ds_conv = nn.Conv2d(dim, dim, 7, padding=3, group=dim, pad_mode="pad") #ds_conv: 深度可分离卷积（Depthwise Separable Convolution），使用7x7的卷积核，padding设为3以保持输出尺寸不变，组数等于特征通道数 dim，这有助于减少参数量并增加模型的效率
        self.net = nn.SequentialCell(
            nn.GroupNorm(1, dim) if norm else nn.Identity(), #分组归一化（GroupNorm，组数为1，意味着逐通道归一化）或者恒等层（如果 norm 参数为 False）。
            nn.Conv2d(dim, dim_out * mult, 3, padding=1, pad_mode="pad"),
            nn.GELU(),
            nn.GroupNorm(1, dim_out * mult),
            nn.Conv2d(dim_out * mult, dim_out, 3, padding=1, pad_mode="pad"),
        )

        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()  #res_conv: 如果输入和输出的特征维度不一致，使用1x1卷积调整维度；如果一致，则直接使用恒等层传递输

    def construct(self, x, time_emb=None):
        h = self.ds_conv(x)
        if exists(self.mlp) and exists(time_emb):
            assert exists(time_emb), "time embedding must be passed in"
            condition = self.mlp(time_emb) #条件调整（如果有时间嵌入）: 如果 time_emb 存在且 mlp 被定义，通过 mlp 处理时间嵌入，并将其扩展成与 h 相同的空间维度后相加，以便将全局时间信息融入局部特征中。
            condition = condition.expand_dims(-1).expand_dims(-1)
            h = h + condition

        h = self.net(h) #特征变换 (net): 将经过条件调整的 h 送入后续的网络层 self.net(h)，这包括归一化、卷积、非线性激活等操作，目的是进一步提炼和丰富特征表示。
        return h + self.res_conv(x) #残差连接: 最后，将经过变换的特征 h 与原始输入 x 经过 res_conv 调整后的版本相加，实现残差连接，这有助于训练深层网络并加速收敛。
        
        
        
ConvNextBlock<
  (ds_conv): Conv2d<input_channels=2, output_channels=2, kernel_size=(7, 7), stride=(1, 1), pad_mode=pad, padding=3, dilation=(1, 1), group=2, has_bias=False, weight_init=<mindspore.common.initializer.HeUniform object at 0x7f92f4bb9730>, bias_init=None, format=NCHW>
  (net): SequentialCell<
    (0): GroupNorm<num_groups=1, num_channels=2>
    (1): Conv2d<input_channels=2, output_channels=4, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=False, weight_init=<mindspore.common.initializer.HeUniform object at 0x7f92f4bb9ca0>, bias_init=None, format=NCHW>
    (2): GELU<>
    (3): GroupNorm<num_groups=1, num_channels=4>
    (4): Conv2d<input_channels=4, output_channels=2, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=False, weight_init=<mindspore.common.initializer.HeUniform object at 0x7f92f4bb9ac0>, bias_init=None, format=NCHW>
    >
  (res_conv): Identity<>
  >

import mindspore as ms
from mindspore import nn, ops
from mindspore.common.initializer import Normal

def exists(x):
    return x is not None

def init_weights(module):
    """权重初始化函数，这里仅作为示例，具体初始化方式根据需求设定"""
    if isinstance(module, nn.Conv2d):
        module.weight.set_data(Normal(0, 0.02)(module.weight.shape))
        if module.bias is not None:
            module.bias.set_data(ms.numpy.zeros(module.bias.shape))

def ds_conv(x, dim):
    """深度可分离卷积操作"""
    return nn.Conv2d(dim, dim, kernel_size=7, padding=3, group=dim, pad_mode="pad")(x)

def mlp(time_emb, dim, time_emb_dim):
    """多层感知机操作，用于时间嵌入处理"""
    if exists(time_emb) and exists(time_emb_dim):
        gelu = nn.GELU()(time_emb)
        dense = nn.Dense(time_emb_dim, dim)(gelu)
        return dense.expand_dims(-1).expand_dims(-1)
    return None

def net(x, dim, dim_out, mult=2, norm=True):
    """主要网络操作"""
    norm_layer = nn.GroupNorm(1, dim) if norm else nn.Identity()
    conv1 = nn.Conv2d(dim, dim_out * mult, kernel_size=3, padding=1, pad_mode="pad")(norm_layer(x))
    gelu = nn.GELU()(conv1)
    norm2 = nn.GroupNorm(1, dim_out * mult)(gelu)
    conv2 = nn.Conv2d(dim_out * mult, dim_out, kernel_size=3, padding=1, pad_mode="pad")(norm2)
    return conv2

def res_connection(x, dim, dim_out):
    """残差连接操作"""
    return nn.Conv2d(dim, dim_out, kernel_size=1)(x) if dim != dim_out else x

def convnext_block_forward(x, time_emb=None, dim=64, dim_out=128, time_emb_dim=False, mult=2, norm=True):
    """
    ConvNextBlock的前向传播函数
    """
    # 深度可分离卷积
    h = ds_conv(x, dim)
    
    # 时间嵌入处理
    condition = mlp(time_emb, dim, time_emb_dim)
    if condition is not None:
        h = h + condition
    
    # 主网络操作
    h = net(h, dim, dim_out, mult, norm)
    
    # 残差连接
    return h + res_connection(x, dim, dim_out)

# 示例使用
dim = 64
dim_out = 128
time_emb_dim = 32
x = ms.Tensor(shape=(1, dim, 64, 64), dtype=ms.float32, init=Normal())  # 假设输入数据形状tensor2 = mindspore.Tensor(shape=(2, 2), dtype=mindspore.float32, init=Normal())
time_emb = ms.Tensor(shape=(1, time_emb_dim), dtype=ms.float32, init=Normal())  # 假设时间嵌入数据形状

output = convnext_block_forward(x, time_emb, dim, dim_out, time_emb_dim)
print("Output shape:", output.shape)

Attention模块

heads=4
dim_head=32
dim_head ** -0.5
hidden_dim = dim_head * heads
hidden_dim

to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, pad_mode='valid', has_bias=False)
to_qkv

to_out = nn.Conv2d(hidden_dim, dim, 1, pad_mode='valid', has_bias=True)
to_out

在卷积神经网络（CNN）中，Conv2d层通常用于从输入特征图中学习局部空间特征。该层的关键参数之一是group，它控制着卷积操作中输入通道与输出通道之间如何进行分组。这个概念被称为“分组卷积”（Grouped Convolution）。

当你设置group=1时，这意味着所有的输入通道会作为一个整体参与卷积运算，与没有指定分组时的传统卷积操作一致。换句话说，在这种情况下，每个输出通道都会受到所有输入通道的贡献，这是标准卷积操作的行为。每个输出通道的计算是通过将输入通道的所有元素与该输出通道对应的卷积核（filter）做内积得到的。

简而言之，当group=1时，分组卷积退化为普通的卷积，即不进行任何特别的分组，所有的输入通道共同作用于生成每一个输出通道的特征图。

from mindspore import Tensor
import mindspore.ops as ops
def show_input(x, y, z):
    return x, y, z

partial = ops.Partial()
partial_show_input = partial(show_input, Tensor(1))  ##show_input(Tensor(1))会错误
output1 = partial_show_input(Tensor(2), Tensor(3))
print(output1)

output2 = partial_show_input(Tensor(3), Tensor(4))
print(output2)

x = ms.Tensor(shape=(1, dim, 64, 64), dtype=ms.float32, init=Normal())
x.shape
print(x)

b, _, h, w = x.shape
to_qkv(x).shape

#to_qkv(x).chunk(3, 1)  ##分三块 （tesor1 tensor2 tensor3）

def rearrange(head, inputs):
    b, hc, x, y = inputs.shape
    c = hc // head
    return inputs.reshape((b, head, c, x * y))

qkv=to_qkv(x).chunk(3, 1)
#qkv  1  128   64 64 
qkv[1]
b, hc, x, y =qkv[1].shape
hc
c = hc // heads
c
qkv[1].reshape(b, heads, c, x * y)

partial1(rearrange, heads) # heads=4

map1 = ops.Map()
partial1=ops.Partial()
map1(partial1(rearrange, heads), qkv)  ####其实就是先传入head 然后传入inputs=qkv进行计算

v=qkv[1].reshape(b, heads, c, x * y)
q=k=qkv[1].reshape(b, heads, c, x * y)
q.shape

sim =ops.bmm(q.swapaxes(2, 3), k)
sim

attn = ops.softmax(sim, axis=-1)
attn

out = ops.bmm(attn, v.swapaxes(2, 3))
out
b, -1, h, w
out.swapaxes(-1, -2)  ###交换维度

out.swapaxes(-1, -2).reshape((1, -1, 64, 64))

这段代码使用了Python中的NumPy库或类似库（如PyTorch、Pandas等数组操作库）中的方法来操作多维数组（通常是张量）。这里是逐步解析这段代码的作用：

1. **out**: 这是一个多维数组（张量），我们不知道它的具体形状，但可以根据后续操作推断出它至少有三个维度。

2. **swapaxes(-1, -2)**: 这个方法用于交换数组中两个轴的位置。在这里，`-1` 和 `-2` 分别代表数组的最后一个轴和倒数第二个轴。在大多数情况下，对于一个三维张量，这可能意味着交换最后一维（通道维度，例如RGB图像的颜色通道）和倒数第二维（高度或宽度之一）。此操作改变了数据的组织方式，但不改变其内容。

3. **reshape((1, -1, 64, 64))**: 这个方法用来改变数组的形状。参数是一个元组，指定了新的形状。在这个例子中：
   - `1`: 表示要在第一个维度创建一个新的批量维度（batch dimension），即使原始张量只有一个样本，也会将其明确表示出来。
   - `-1`: 是一个特殊的值，意味着这一维度的大小应该自动计算，以便使总的元素数量保持不变。这里，它会根据原张量的总元素数量以及已经指定的其他维度大小自动确定这一维度的值。
   - `(64, 64)`: 明确指定了重塑后张量的最后两个维度应为64x64，这通常意味着在空间上调整为64像素宽和高。

综上所述，这段代码首先交换了原张量的最后两维，然后将其重塑为一个新的四维张量，其中包含一个批量（大小为1），一个自动计算大小的维度（对应于通道数或某些其他分组），以及64x64的空间维度。这样的操作常见于深度学习中，用于调整张量的形状以适应特定层的输入要求，或者为了数据处理和可视化的目的。

out = out.swapaxes(-1, -2).reshape((b, -1, h, w))
out

x = ms.Tensor(shape=(1, dim, 64, 64), dtype=ms.float32, init=Normal())
x.shape
print(x)

x.var(1, keepdims=True)
在深度学习或数据分析的上下文中，`x.var(1, keepdims=True)` 这段代码是用来计算张量 `x` 沿着指定维度（本例中是第1维度）的方差，并且保持输出的维度与原始张量相同，除了计算方差的那一维被缩减为1。下面是对这段代码的详细解析：

- **x.var()**: 这里 `var` 是一个方法，用于计算张量（在诸如PyTorch、NumPy等库中）的方差。方差是衡量数据分散程度的一个统计量，计算公式为各个数据与其平均值之差的平方的平均值。

- **1**: 这个参数指定了沿着哪个维度进行方差计算。在大多数库中，维度计数是从0开始的，所以 `1` 表示沿着第二个维度进行操作。例如，如果你有一个形状为`(batch_size, channels, height, width)`的四维张量，那么计算就是沿着通道维度进行的。

- **keepdims=True**: 这是一个布尔标志，用来指示在计算完成后是否保持输出的维度与输入相同。当设置为 `True` 时，即使计算过程中某个维度被缩减（因为计算了该维度上的统计量，如均值、方差等），输出的张量仍然会保持与输入相同的维度数量，被计算的维度大小会被设为1。这样做有利于后续的广播操作，使得计算结果可以直接与原始张量进行维度匹配，进行进一步的计算或操作，而无需手动调整形状。

综上所述，`x.var(1, keepdims=True)` 的作用是计算张量 `x` 在其第二个维度上的方差，并且保持输出张量的维度结构不变，即计算方差的维度大小变为1，其他维度保持原样。这对于需要在计算统计量的同时维持张量维度一致性，以便进行后续操作的场景非常有用。

# 使用Normal分布初始化一个Tensor
x = ms.Tensor(shape=(1, 64, 4, 4), dtype=ms.float32, init=Normal(sigma=0.5))
print(x)

b, _, h, w = x.shape
h
qkv = to_qkv(x).chunk(3, 1)
q, k, v = map1(partial1(rearrange, heads), qkv)
q

ops.softmax(q, -2)

x = ms.Tensor(shape=(1,2, 4, 4), dtype=ms.float32, init=Normal(sigma=0.5))
print(x)
ops.softmax(x, -2) #，这意味着所有元素都会被转换为0到1之间的值，并且张量中所有元素的和为1。

ops.sum(ops.softmax(x, -2), -2) 

print("Original Tensor Slice:\n", ops.softmax(x, -2) [0, 0, 0,:])  # 打印第一个批次的前3个通道的元素

ops.softmax(x, -2) [0, 0, 0,:].sum()

b, _, h, w = x.shape
h,w,v.shape

class Attention(nn.Cell):
    def __init__(self, dim, heads=4, dim_head=32):
        super().__init__()
        self.scale = dim_head ** -0.5
        self.heads = heads
        hidden_dim = dim_head * heads

        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, pad_mode='valid', has_bias=False)
        self.to_out = nn.Conv2d(hidden_dim, dim, 1, pad_mode='valid', has_bias=True)
        self.map = ops.Map()
        self.partial = ops.Partial()

    def construct(self, x):
        b, _, h, w = x.shape
        qkv = self.to_qkv(x).chunk(3, 1)
        q, k, v = self.map(self.partial(rearrange, self.heads), qkv)  ####就是rearrange的两个参数  得到 tuple 有三个元素都是tensor类型

        q = q * self.scale

        # 'b h d i, b h d j -> b h i j'    
        sim = ops.bmm(q.swapaxes(2, 3), k)   ##q  shape=[1, 4, 32, 4096]
        attn = ops.softmax(sim, axis=-1) 
        # 'b h i j, b h d j -> b h i d'
        out = ops.bmm(attn, v.swapaxes(2, 3))
        out = out.swapaxes(-1, -2).reshape((b, -1, h, w))

        return self.to_out(out)

class LayerNorm(nn.Cell):
    def __init__(self, dim):
        super().__init__()
        self.g = Parameter(initializer('ones', (1, dim, 1, 1)), name='g')

    def construct(self, x):
        eps = 1e-5
        var = x.var(1, keepdims=True)
        mean = x.mean(1, keep_dims=True)
        return (x - mean) * rsqrt((var + eps)) * self.g

class LinearAttention(nn.Cell):
    def __init__(self, dim, heads=4, dim_head=32):
        super().__init__()
        self.scale = dim_head ** -0.5
        self.heads = heads
        hidden_dim = dim_head * heads
        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, pad_mode='valid', has_bias=False)

        self.to_out = nn.SequentialCell(
            nn.Conv2d(hidden_dim, dim, 1, pad_mode='valid', has_bias=True),
            LayerNorm(dim)
        )

        self.map = ops.Map()
        self.partial = ops.Partial()

    def construct(self, x):
        b, _, h, w = x.shape
        qkv = self.to_qkv(x).chunk(3, 1)
        q, k, v = self.map(self.partial(rearrange, self.heads), qkv)

        q = ops.softmax(q, -2)
        k = ops.softmax(k, -1)

        q = q * self.scale
        v = v / (h * w)

        # 'b h d n, b h e n -> b h d e'
        context = ops.bmm(k, v.swapaxes(2, 3))
        # 'b h d e, b h d n -> b h e n'
        out = ops.bmm(context.swapaxes(2, 3), q)

        out = out.reshape((b, -1, h, w))
        return self.to_out(out)
        
        
这两个注意力类，分别是标准的多头自注意力（`Attention`）和线性注意力（`LinearAttention`）的实现，它们在计算注意力权重的方式上有所差异，导致了不同的性能特性和适用场景。下面是它们各自计算公式的详细解析：

### Attention 类（多头自注意力）

1. **计算Query, Key, Value**:
   - 输入张量 `x` 经过线性变换（`to_qkv`），然后切分得到 Query(`q`), Key(`k`), Value(`v`)，每个被分为 `heads` 个头。
   - 对 `q` 应用缩放因子 `\(\sqrt{dim\_head}\)`。

2. **注意力分数计算**:
   - 使用矩阵乘法（`ops.bmm`）计算 `q` 和 `k` 之间的点积相似度，得到注意力分数矩阵 `sim`，形状为 `[b, h, i, j]`，其中 `i` 和 `j` 分别代表查询和键的序列长度（在此上下文中，实际上是空间维度的展开，如图像的宽和高）。

3. **应用softmax得到注意力权重**:
   - 对 `sim` 沿着最后一个维度（即 `-1`，代表 `j`）应用softmax函数，得到注意力权重分布。

4. **加权求和得到输出**:
   - 使用注意力权重矩阵乘以 `v`（`ops.bmm(attn, v.swapaxes(2, 3))`），然后通过一系列变换（交换和重塑）恢复原始形状。

### LinearAttention 类（线性注意力）

1. **计算Query, Key, Value**:
   - 同样，输入 `x` 经过线性变换切分成 `q`, `k`, `v`，每部分分为 `heads` 个头。

2. **软对齐操作**:
   - 直接对 `q` 和 `k` 应用softmax，分别沿 `-2`（查询的特征维度）和 `-1`（键的序列维度）进行归一化。这一步与标准自注意力中的操作不同，它直接对查询和键进行软对齐，而不是计算它们之间的点积相似度。

3. **计算上下文向量**:
   - 将归一化的 `k` 与归一化的 `v` 通过矩阵乘法得到上下文向量 `context`。这一步骤相当于在键值对之间进行了加权求和，但权重是由 `k` 和 `v` 自身的分布决定的，而非基于点积的相似度。

4. **聚合并输出**:
   - 将 `context` 与 `q` 通过矩阵乘法得到最终的输出 `out`，再经过形状调整和线性变换（包括LayerNorm）后返回。

**总结**:
- 标准的多头自注意力（`Attention`）通过计算query和key之间的点积来衡量相似度，然后基于这个相似度进行加权求和得到value的输出，这适用于需要捕获长距离依赖和复杂模式的场景。
- 线性注意力（`LinearAttention`）则通过对query和key分别进行softmax操作，简化了注意力分数的计算，降低了计算复杂度，适合资源受限或需要快速计算的场景，但可能牺牲了一些模型的表达能力，因为它不直接衡量query和key之间的匹配度。

`nn.GroupNorm(1, dim)` 这段代码在深度学习框架（如PyTorch或MindSpore）中创建了一个组归一化（Group Normalization）层。组归一化是一种正则化技术，旨在解决Batch Normalization在小批量大小或动态批量大小情况下的局限性。它将输入通道划分为若干组，并在组内执行标准化处理，以减少内部协变量偏移（internal covariate shift）问题，从而加速训练过程并提高模型的泛化能力。

下面是对`nn.GroupNorm(1, dim)`参数的详细解析及其计算效果：

- **1**: 这是`num_groups`参数，表示通道分组的数量。在本例中，设置为1意味着每个组只包含1个通道。这个特殊的设置使得GroupNorm退化为层归一化（Layer Normalization），因为它对每个通道独立地执行归一化操作。层归一化通常用于处理RNNs或没有空间结构的数据（如序列数据），因为它不依赖于数据的宽度或高度。

- **dim**: 这是输入张量的特征维度，也就是通道数。在卷积神经网络中，特征维度通常对应于通道数。GroupNorm需要知道通道数来正确地划分通道组并进行归一化处理。

**计算效果**：

对于输入张量`X`，形状通常为`(N, C, H, W)`，其中`N`是批量大小，`C`是通道数（特征维度），`H`和`W`分别是特征图的高度和宽度。当使用`nn.GroupNorm(1, dim)`时：

1. **通道分组**：由于`num_groups`设置为1，每个通道被单独视为一组，因此实际上不会进行真正的分组操作，每个通道独立处理。

2. **计算均值和方差**：对于每个通道，计算该通道所有元素（跨越所有批量、高度和宽度）的平均值（均值）和标准差（方差）。

3. **归一化**：对于每个通道内的每个元素，使用如下公式进行归一化：
   - 先减去该通道的均值，然后除以该通道的方差（加上一个很小的常数epsilon，以避免除以零的问题）。

4. **缩放和平移**：之后，通常会有一个可学习的缩放系数（gamma）和偏移量（beta），用于恢复归一化后的数据的表达能力和灵活性。归一化后的值乘以`gamma`并加上`beta`，得到最终的输出。

因此，`nn.GroupNorm(1, dim)`的主要效果是，对输入数据在每个通道上独立执行归一化，调整和稳定每个特征通道的分布，使得模型对不同层间输入的变化更加鲁棒，有助于训练稳定性和模型性能。当组数为1时，其行为等价于Layer Normalization，常用于处理序列数据或不规则结构的数据，以及小批量训练场景。

        
        ### 条件U-Net

我们已经定义了所有的构建块（位置嵌入、ResNet/ConvNeXT块、Attention和组归一化），现在需要定义整个神经网络了。请记住，网络 $\mathbf{\epsilon}_\theta(\mathbf{x}_t, t)$ 的工作是接收一批噪声图像+噪声水平，并输出添加到输入中的噪声。

更具体的：
网络获取了一批`(batch_size, num_channels, height, width)`形状的噪声图像和一批`(batch_size, 1)`形状的噪音水平作为输入，并返回`(batch_size, num_channels, height, width)`形状的张量。

网络构建过程如下：

- 首先，将卷积层应用于噪声图像批上，并计算噪声水平的位置

- 接下来，应用一系列下采样级。每个下采样阶段由2个ResNet/ConvNeXT块 + groupnorm + attention + 残差连接 + 一个下采样操作组成

- 在网络的中间，再次应用ResNet或ConvNeXT块，并与attention交织

- 接下来，应用一系列上采样级。每个上采样级由2个ResNet/ConvNeXT块+ groupnorm + attention + 残差连接 + 一个上采样操作组成

- 最后，应用ResNet/ConvNeXT块，然后应用卷积层

最终，神经网络将层堆叠起来，就像它们是乐高积木一样（但重要的是[了解它们是如何工作的](http://karpathy.github.io/2019/04/25/recipe/)）。
        

class Unet(nn.Cell):
    def __init__(
            self,
            dim,
            init_dim=None,
            out_dim=None,
            dim_mults=(1, 2, 4, 8),
            channels=3,
            with_time_emb=True,
            convnext_mult=2,
    ):
        super().__init__()

        self.channels = channels

        init_dim = default(init_dim, dim // 3 * 2)
        self.init_conv = nn.Conv2d(channels, init_dim, 7, padding=3, pad_mode="pad", has_bias=True)

        dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
        in_out = list(zip(dims[:-1], dims[1:]))

        block_klass = partial(ConvNextBlock, mult=convnext_mult)

        if with_time_emb:
            time_dim = dim * 4
            self.time_mlp = nn.SequentialCell(
                SinusoidalPositionEmbeddings(dim),
                nn.Dense(dim, time_dim),
                nn.GELU(),
                nn.Dense(time_dim, time_dim),
            )
        else:
            time_dim = None
            self.time_mlp = None

        self.downs = nn.CellList([])
        self.ups = nn.CellList([])
        num_resolutions = len(in_out)

        for ind, (dim_in, dim_out) in enumerate(in_out):
            is_last = ind >= (num_resolutions - 1)

            self.downs.append(
                nn.CellList(
                    [
                        block_klass(dim_in, dim_out, time_emb_dim=time_dim),
                        block_klass(dim_out, dim_out, time_emb_dim=time_dim),
                        Residual(PreNorm(dim_out, LinearAttention(dim_out))),
                        Downsample(dim_out) if not is_last else nn.Identity(),
                    ]
                )
            )

        mid_dim = dims[-1]
        self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)
        self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim)))
        self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)

        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
            is_last = ind >= (num_resolutions - 1)

            self.ups.append(
                nn.CellList(
                    [
                        block_klass(dim_out * 2, dim_in, time_emb_dim=time_dim),
                        block_klass(dim_in, dim_in, time_emb_dim=time_dim),
                        Residual(PreNorm(dim_in, LinearAttention(dim_in))),
                        Upsample(dim_in) if not is_last else nn.Identity(),
                    ]
                )
            )

        out_dim = default(out_dim, channels)
        self.final_conv = nn.SequentialCell(
            block_klass(dim, dim), nn.Conv2d(dim, out_dim, 1)
        )

    def construct(self, x, time):
        x = self.init_conv(x)

        t = self.time_mlp(time) if exists(self.time_mlp) else None

        h = []

        for block1, block2, attn, downsample in self.downs:
            x = block1(x, t)
            x = block2(x, t)
            x = attn(x)
            h.append(x)

            x = downsample(x)

        x = self.mid_block1(x, t)
        x = self.mid_attn(x)
        x = self.mid_block2(x, t)

        len_h = len(h) - 1
        for block1, block2, attn, upsample in self.ups:
            x = ops.concat((x, h[len_h]), 1)
            len_h -= 1
            x = block1(x, t)
            x = block2(x, t)
            x = attn(x)

            x = upsample(x)
        return self.final_conv(x)  修改为函数式编程

这段代码定义了一个基于PyTorch或类似框架（这里看起来像是MindSpore，因为使用了`nn.Cell`）的U-Net架构，该架构常用于图像生成、图像分割等任务，尤其是在处理具有时间序列信息的任务时（如在给定时间步长的情况下生成图像）。此实现包含一些特定的组件，比如使用`ConvNextBlock`作为基本构建单元，以及线性注意力和多尺度结构，使得网络能够在不同分辨率下捕获并整合特征。下面是对代码的逐部分解析：

### 初始化
- `dim` 是网络的基本维度，控制了网络的宽度。
- `init_dim` 初始化卷积层的输出维度，默认为 `dim // 3 * 2`。
- `out_dim` 输出卷积层的维度，默认与输入通道数相同。
- `dim_mults` 控制每一层下采样/上采样的通道数乘数，决定网络的宽度如何随着深度增加而增加。
- `channels` 输入图像的通道数，默认为3，对应RGB图像。
- `with_time_emb` 是否包含时间嵌入，这对于时序数据很重要。
- `convnext_mult` 控制`ConvNextBlock`中乘法因子的超参数。

### 层定义
- **初始化卷积 (`init_conv`)**：对输入图像进行初步特征提取。
- **时间嵌入 (`time_mlp`)**：如果启用，通过正弦位置嵌入和多层感知机来编码时间信息，增强模型的时间意识。
- **下采样路径 (`downs`)**：由多个模块组成，包括两个`ConvNextBlock`（用于特征提取）、线性注意力模块（通过`LinearAttention`实现）以及下采样操作（通过`Downsample`），用于逐步减小空间尺寸同时增加特征维度。
- **中间块 (`mid_block1`, `mid_attn`, `mid_block2`)**：在下采样和上采样路径之间，用于处理最精细的特征层次，包含两个`ConvNextBlock`、一个注意力模块，帮助模型在不同特征尺度间建立联系。
- **上采样路径 (`ups`)**：与下采样路径相反，逐步增加空间尺寸并减少特征维度，通过上采样操作（`Upsample`）和跳跃连接（将对应下采样层的特征拼接到当前层）实现。
- **最终卷积 (`final_conv`)**：将特征映射回所需的输出通道数，通常与输入图像的通道数相同。

### 前向传播 (`construct` 方法)
1. **初始化卷积**：对输入图像`x`应用初始卷积。
2. **时间嵌入处理**：如果启用，根据当前时间`t`计算时间嵌入。
3. **下采样过程**：依次遍历每个下采样模块，应用卷积块、注意力模块和下采样操作，并保存每一层的输出用于后续的跳跃连接。
4. **中间处理**：在下采样和上采样之间，通过一系列操作处理最深层的特征。
5. **上采样过程**：与下采样过程相反，结合跳跃连接的特征，逐步上采样并减少特征维度。
6. **输出**：最后通过一个`ConvNextBlock`和一个卷积层得到最终输出。

这个U-Net架构结合了`ConvNextBlock`、线性注意力和多尺度特征处理，旨在提高图像生成任务的表现，尤其是当时间序列信息对于任务至关重要时。

image_size = 28
channels = 1
batch_size = 2

image_size, batch_size=16, channels=3

# 定义 Unet模型
unet_model = Unet(
    dim=image_size,
    channels=channels,
    dim_mults=(1, 2, 4,)
)

#x = ms.Tensor(shape=(1, dim, 64, 64), dtype=ms.float32, init=Normal())
#x.shape
#print(x)

import mindspore as ms
from mindspore.common.initializer import Normal

# 参数定义
image_side_length = 32  # 图像的宽和高的像素数
channels = 3  # 图像通道数，这里假设处理的是RGB图像
batch_size = 2  # 批次大小

# 定义 Unet模型
# 注意：此处的dim应该根据模型设计具体指定，但基于您的代码，我们保持原样
unet_model = Unet(dim=image_side_length, channels=channels, dim_mults=(1, 2, 4,))

# 构建输入数据
x = ms.Tensor(shape=(batch_size, channels, image_side_length, image_side_length), dtype=ms.float32, init=Normal())
x.shape  # 显示数据形状
print(x)  # 打印数据（显示初始化后的随机值）