StyleGAN-v2 ADA 的 pytorch 代码分析与实际使用

$\mathrm{R}\mathrm{e}\mathrm{s}\mathrm{o}\mathrm{u}\mathrm{r}\mathrm{c}\mathrm{e}\mathrm{s}$

 📋 $\mathrm{P}\mathrm{a}\mathrm{p}\mathrm{e}\mathrm{r}$ >> $\mathrm{T}\mathrm{r}\mathrm{a}\mathrm{i}\mathrm{n}\mathrm{i}\mathrm{n}\mathrm{g}\mathrm{G}\mathrm{e}\mathrm{n}\mathrm{e}\mathrm{r}\mathrm{a}\mathrm{t}\mathrm{i}\mathrm{v}\mathrm{e}\mathrm{A}\mathrm{d}\mathrm{v}\mathrm{e}\mathrm{r}\mathrm{s}\mathrm{a}\mathrm{r}\mathrm{i}\mathrm{a}\mathrm{l}\mathrm{N}\mathrm{e}\mathrm{t}\mathrm{w}\mathrm{o}\mathrm{r}\mathrm{k}\mathrm{s}\mathrm{w}\mathrm{i}\mathrm{t}\mathrm{h}\mathrm{L}\mathrm{i}\mathrm{m}\mathrm{i}\mathrm{t}\mathrm{e}\mathrm{d}\mathrm{D}\mathrm{a}\mathrm{t}\mathrm{a}$
💻 $\mathrm{C}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{s}$ >>tensorflow/pytorch
📰 $\mathrm{B}\mathrm{l}\mathrm{o}\mathrm{g}\mathrm{s}$  >> $\mathrm{C}\mathrm{S}\mathrm{D}\mathrm{N}:$ $\mathrm{T}\mathrm{r}\mathrm{a}\mathrm{i}\mathrm{n}\mathrm{i}\mathrm{n}\mathrm{g}\mathrm{G}\mathrm{e}\mathrm{n}\mathrm{e}\mathrm{r}\mathrm{a}\mathrm{t}\mathrm{i}\mathrm{v}\mathrm{e}\mathrm{A}\mathrm{d}\mathrm{v}\mathrm{e}\mathrm{r}\mathrm{s}\mathrm{a}\mathrm{r}\mathrm{i}\mathrm{a}\mathrm{l}\mathrm{N}\mathrm{e}\mathrm{t}\mathrm{w}\mathrm{o}\mathrm{r}\mathrm{k}\mathrm{s}\mathrm{w}\mathrm{i}\mathrm{t}\mathrm{h}\mathrm{L}\mathrm{i}\mathrm{m}\mathrm{i}\mathrm{t}\mathrm{e}\mathrm{d}\mathrm{D}\mathrm{a}\mathrm{t}\mathrm{a}$

$\mathrm{A}\mathrm{n}\mathrm{a}\mathrm{l}\mathrm{y}\mathrm{s}\mathrm{i}\mathrm{s}$

$1.$ 计算 $L_{adv}$ 时，前馈真实样本或伪样本给鉴别器 $D$;

## 对生成伪样本
fake_logits = self.run_D(fake_img, fake_c)
## 对真实训练样本
real_logits = self.run_D(real_img, real_c)

$2.$ run_D 的具体细节为：

def run_D(self, img, c):
    ## 使用定义的含 p 的数据增强流水线作 aug
    img = self.augment_pipe(img) 
    logits = self.D(img, c)
    return logits

$3.$ 记录 $L_{adv}$ 和 $\mathbb{E}[\mathrm{s}\mathrm{i}\mathrm{g}\mathrm{n}({\mathrm{D}}_{\mathrm{t}\mathrm{r}\mathrm{a}\mathrm{i}\mathrm{n}})]$ （这里使用的是 WGAN，所以 $L_{adv}$ 计算方式比较简单，$\min /\max$ ?_logits 即可）

training_stats.report('Loss/scores/real', real_logits)
training_stats.report('Loss/signs/real' , real_logits.sign()) ## 👈 

$4.$ 具体的，每一个被 $\mathrm{r}\mathrm{e}\mathrm{p}\mathrm{o}\mathrm{r}\mathrm{t}\mathrm{e}\mathrm{d}$ 的状态（统计数据 $\mathrm{s}\mathrm{t}\mathrm{a}\mathrm{t}\mathrm{i}\mathrm{s}\mathrm{t}\mathrm{i}\mathrm{c}$）被记录了 3 个统计量：

## `elems` 是一个形参，这里考虑 tensor `logits`
moments = torch.stack([
    torch.ones_like(elems).sum(), ## 记录数量 (count)
    elems.sum(),                  ## 求和
    ## 计算 E[sign(D_train)] 只需要前两个统计量，即：moments[1]/moments[0]
    elems.square().sum(),
])

$5.$ 累计前面 $4.$ 记录的数据 moments of real_logits，累积（会使用到 _moments.add_(moments)） $N=4$ 次迭代（$\mathrm{i}\mathrm{t}\mathrm{e}\mathrm{r}/\mathrm{m}\mathrm{i}\mathrm{n}\mathrm{i}\mathrm{b}\mathrm{a}\mathrm{t}\mathrm{c}\mathrm{h}$）
$6.$ 通过获取 $\mathbb{E}[\mathrm{s}\mathrm{i}\mathrm{g}\mathrm{n}({\mathrm{D}}_{\mathrm{t}\mathrm{r}\mathrm{a}\mathrm{i}\mathrm{n}})]$ 来动态更新 p

# Execute ADA heuristic.
if (ada_stats is not None) and \ ## 是否使用 ADA 这一项技术
   (batch_idx % ada_interval == 0): ## N 值
    ada_stats.update()
    adjust = np.sign(ada_stats['Loss/signs/real'] - ada_target) \ ## ada_target 是 r_t 的阈值，文中设置是 0.6
                                                                  ## ada_stats['Loss/signs/real'] = moments_of_real_logits[1]/moments_of_real_logits[0]
             * (batch_size * ada_interval) / (ada_kimg * 1000)    ## 增益，the gain := (BxN)/SCALE, `B` is batch size, `N` is # of batches; all-in-all, it is FIXED.
    ## 更新 p 值
    augment_pipe.p.copy_((augment_pipe.p + adjust)\               ## D 偏强，则 adjust 为正，Aug 强度适当增大；D 偏弱，则 adjust 为负，Aug 强度适当减弱
                         .max(misc.constant(0, device=device)))   ## clip/truncate，限制概率在有效范围

$\mathrm{U}\mathrm{s}\mathrm{a}\mathrm{g}\mathrm{e}\mathrm{o}\mathrm{f}\mathrm{A}\mathrm{u}\mathrm{g}\mathrm{P}\mathrm{i}\mathrm{p}\mathrm{e}$

引用脚本文件 ./training/augment.py，直接初始化 nn.Module 模块——

import augment

aug_pipe = augment.AugmentPipe()
## input --type=torch.tensor --size=(N,C,H,W)
aug_input = aug_pipe(input)

$\mathrm{B}\mathrm{T}\mathrm{W}$，这个项目对于 pytorch 多线程、多卡并行训练编程有非常好的借鉴性，木奉👍！