主要代码将会完全来自https://github.com/lucidrains/denoising-diffusion-pytorch/tree/main
Norm 需要知道dim 来训练ｇ
PreNorm → RMSNorm 后再通过fn
Residual block 只需传入func 之后通过forward
主要组成部分-1 Positional embedding
Frequency变短 Period(T = 1/w) 就是分数下面部分随dim变长
这里前两行emb计算的就是上边w 也就是frequency 这里使用了exponential形式而已
接下来对于x([1,2,3…Sequence length])的每个t
这里的x应该是t 是每次UNET forward传入值 size为(batch_size,) 数值统一为time_step 这样对于每batch训练有共同emb emb会之后传入mlp
都生成了对应于dim位置i 的总长为half_dim
最终 前半dim为sin 后半为cos
原方法应该是交替的 这里暂时不知道原因
接下来的方法 其他都一样 只是freq换成了 learned parameter
值得注意的是这里的 x[:,None] 替换成了 einops(tensor operation库).rearrange(x, ‘b → b 1’) 都代表将(b,)转换为(b, 1)
主要组成部分-2 Conv Block
Projection(Conv 升降维 特征提取) → RMSNorm → Scale shift → Activation → Dropout
Block1(With time embedding added as scale shift)
-
Block2
-
Original Copy of x
这里的scale shift很重要 对于每个time stamp会有不同的activation 使得模型知道处于哪一部denoising之中 提供time-related information
这里的每一个batch 对应同一time step的多张图片 这些图片的每个channel共享一个(time_emb_dim)的参数 ← 从之前pos_emb的(T * time_emb_dim) 中对应取出
之后扩大到(batch_size, time_emb_dim) → 通过mlp转换为(batch_size, dim*2) 图像的dim对应channel
再之后分成两份 线性添加在activation之前
主要组成部分-３ Attention（Linear & 普通）
hidden_dim = heads * head_dim
sequence length 放最后因为 图片训练直接使用pixel value 不像文字拥有embedding 使得每个word vector为n*d 有后续维度
einops.rearrange(x, ‘a b c → ａ （ｂ ｃ）’)
这里使用lambda map(lambda x: rearrange(x, ‘b (h c) n → b h c n’), qkv)
两个softmax分别作用不同ｑｋ不同dim上
Linear Attention 去掉了Ｑ＊Ｋ的ｎ＾２中间部（Attention score）
而是先计算Context（包括对Ｑ，Ｋ的Kernal function） 这里使用Softmax 目的? Maintain linear scaling
主要组成部分-4 UNet
标注每个layer的dimension 这里dims感觉语法复杂了 in_out写法值得学习
initial_dim是传入的 之后dim都是scale的(通过dim_mults)
self.condition会将之前output重新带入训练 input channel * 2 但是也会线性转换 input_channel → input_dim
创建time_step embedding
dim为传入值 learned_sinusoidal_dim设定为16
partial freezes argument 下次创建时可以省略这两个
2 * Resnet block(dim不变)
-
Residual(Norm后self linear attention, x)
-
Down sample(change dim)
UNET 中间步
Resnet block1
-
Residual(Norm后self linear attention, x)
-
Resnet block2
UP时需要注意的是 Resnet_block的dims是(dim_out + dim_in, dim_out)
最后经过final_res_block（init_dim * 2, init_dim）(？) 和一个Conv1d
x_self_cond是选择传入的选项 但是即使不加 依旧会torch.cat上torch.zeros_like 跟x大小一样的0
取决于x_self_cond
init_conv(Conv1d) 将channel转换为init_dim 视野较大 size=7
Module_list 使得forward可以像zip一样提取每一个运算
h 记住每一层的每个大运算结果 添加到UNET右边 作为input