0
0

使用Modulated Convolutions修改 StarGAN V2

在本教程中，我们将替换 StarGAN V2 模型中的自适应实例归一化（AdaIN）层，并在分辨率为 512x512 像素的图像上对其进行训练。

为什么是 StarGAN V2？

今天，有许多模型可以生成高质量的图像。具体来说，对于属性交换任务（2021 年），最好的质量是由 StyleGAN 的进一步发展的模型提供的，或者是通过对其进行提炼而获得的，这需要大量的计算时间来训练新的数据域。在单个 Google Colab GPU 上从头开始训练 24 小时后，所提出的模型会生成文章开头所示的图像。

StarGAN V2 架构

StarGAN V2[1] 是一个图像到图像模型，它使用由条件编码器管理的 AdaIN 层来传输图像风格。它分别使用有关对象结构及其纹理的信息，从而允许用户获得组合图像。

StarGAN 与图像生成相关的部分如下图所示。它们包括 ResNet-like[2] 编码器——用绿色标记，解码器带有 AdainResBlk 模块（将在下面描述）——紫色，以及一组具有共享头层的条件相关样式信息编码器（灰蓝色）——用绿松石标记。

StarGAN 的工作原理如下。一开始，风格编码器从图像中提取低级特征。然后生成器编码对象的几何信息并将其提供给 AdainResBlk 模块的金字塔。

每个 AdainResBlk 块都包含 StyleGAN 的自适应实例归一化（AdaIN）模块 [3]，它通过从样式编码器接收到的信息来调制抽象对象的几何表示。

让我们开始我们的项目，用来自 StyleGAN 2[4] 的调制卷积替换 AdaIN 归一化。

StarGAN 修改

首先，我们需要原始 StarGAN 的 repo：git clone https://github.com/clovaai/stargan-v2.git。

AdainResBlk 的源代码位于 core/model.py 文件中。代码如下所示。

  1.  class AdainResBlk(nn.Module):
  2.      def __init__(self, dim_in, dim_out, style_dim=64, w_hpf=0,
  3.                   actv=nn.LeakyReLU(0.2), upsample=False):
  4.          # ...
  5.  
  6.      def _build_weights(self, dim_in, dim_out, style_dim=64):
  7.          self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
  8.          self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
  9.          self.norm1 = AdaIN(style_dim, dim_in)
  10.          self.norm2 = AdaIN(style_dim, dim_out)
  11.          if self.learned_sc:
  12.              self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
  13.  
  14.      def _shortcut(self, x):
  15.          if self.upsample:
  16.              x = F.interpolate(x, scale_factor=2, mode='nearest')
  17.          if self.learned_sc:
  18.              x = self.conv1x1(x)
  19.          return x
  20.  
  21.      def _residual(self, x, s):
  22.          x = self.norm1(x, s)
  23.          x = self.actv(x)
  24.          if self.upsample:
  25.              x = F.interpolate(x, scale_factor=2, mode='nearest')
  26.          x = self.conv1(x)
  27.          x = self.norm2(x, s)
  28.          x = self.actv(x)
  29.          x = self.conv2(x)
  30.          return x
  31.  
  32.      def forward(self, x, s):
  33.          out = self._residual(x, s)
  34.          if self.w_hpf == 0:
  35.              out = (out + self._shortcut(x)) / math.sqrt(2)
  36.          return out

现在，我们用 lucidrains StyleGAN 2 模块 [5] 替换了 AdainResBlk。类似于 AdainResBlk 的功能在 GeneratorBlock 类（文件 stylegan2_pytorch.py）中实现。让我们将这个类及其依赖项——Conv2DMod、Blur 和 RGBBlock 复制到我们的仓库中。

生成器块的最终版本如下所示。

  1.  from modulated_convolution import Conv2DMod, RGBBlock    
  2.      
  3.  class GenResBlk(nn.Module):
  4.      def __init__(self, dim_in, dim_out, style_dim=64, fade_num_channels=4, fade_num_hidden=32,
  5.                   actv=nn.LeakyReLU(0.2), upsample=False):
  6.          super().__init__()
  7.          self.actv = actv
  8.          self.upsample = upsample
  9.          self.needSkipConvolution = dim_in != dim_out
  10.          self.conv1 = Conv2DMod(dim_in, dim_out, 3, stride=1, dilation=1)
  11.          self.conv2 = Conv2DMod(dim_out, dim_out, 3, stride=1, dilation=1)
  12.          self.style1 = nn.Linear(style_dim, dim_in)
  13.          self.style2 = nn.Linear(style_dim, dim_out)
  14.          if self.needSkipConvolution:
  15.              self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
  16.          self.toRGB = RGBBlock(style_dim, dim_out, upsample, 3)
  17.  
  18.      def forward(self, x, rgb, s):
  19.          if self.upsample:
  20.              x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
  21.          if self.needSkipConvolution:
  22.              x_ = self.conv1x1(x)
  23.          else:
  24.              x_ = x
  25.          s1 = self.style1(s)
  26.          x = self.conv1(x, s1)
  27.          x = self.actv(x)
  28.          s2 = self.style2(s)
  29.          x = self.conv2(x, s2)
  30.          x = self.actv(x + x_)
  31.          rgb = self.toRGB(x, rgb, s)
  32.          return x, rgb

为简单起见，我们不会改变 StyleGAN 使用两个流——特征流和 RGB 图像流的原始概念，因此有必要修改生成器的前向方法。

替换最近的几行：

  1.  def forward(self, x, s, masks=None):
  2.      x = self.from_rgb(x)
  3.      cache = {}
  4.      for block in self.encode:
  5.          if (masks is not None) and (x.size(2) in [32, 64, 128]):
  6.              cache[x.size(2)] = x
  7.          x = block(x)
  8.      for block in self.decode:
  9.          x = block(x, s)
  10.          if (masks is not None) and (x.size(2) in [32, 64, 128]):
  11.              mask = masks[0] if x.size(2) in [32] else masks[1]
  12.              mask = F.interpolate(mask, size=x.size(2), mode='bilinear')
  13.              x = x + self.hpf(mask * cache[x.size(2)])
  14.      return self.to_rgb(x)

使用下一个代码块：

  1.  def forward(self, x, s, masks=None):
  2.      x = self.from_rgb(x)
  3.      cache = {}
  4.      for block in self.encode:
  5.          if (masks is not None) and (x.size(2) in [32, 64, 128]):
  6.              cache[x.size(2)] = x
  7.          x = block(x)
  8.      rgb = None
  9.      for block in self.decode:
  10.          x, rgb = block(x, rgb, s)
  11.          if (masks is not None) and (x.size(2) in [32, 64, 128]):
  12.              mask = masks[0] if x.size(2) in [32] else masks[1]
  13.              mask = F.interpolate(mask, size=x.size(2), mode='bilinear')
  14.              x = x + self.hpf(mask * cache[x.size(2)])
  15.      return rgb

为避免测试调用时出现 OOM，请在 debug_image 函数（文件 utils.py）中注释“潜在引导图像合成”和“参考引导图像合成”块。

对于 512x512 图像的训练模型，我们必须将批量大小减少到 1。为了稳定训练过程，我们将使用假图像缓冲区（来自 pytorch-CycleGAN-and-pix2pix repo），它允许我们使用以下方法更新鉴别器的权重 生成数据的历史记录，而不是最新的假输出。

使用 Colab 的注意事项

如果您将在 Colab 环境中训练模型，您可以修改 _save_checkpoint 和 _load_checkpoint 函数中的步骤参数（在任何情况下，记得Google Drive 创建备份）并在将当前模型复制到 Drive 的训练函数中添加下一行：

  1.  # save model checkpoints
  2.  if (i+1) % args.save_every == 0:
  3.      self._save_checkpoint(step=i+1)
  4.      print("Saving on GDrive...")
  5.      import subprocess
  6.      subprocess.run(f"cp --force -R ./expr/ /content/drive/MyDrive/stargan_animals_expr/", shell=True, check=True)

模型训练

将 AFHQ 放入 data/ 文件夹后，我们就可以开始训练了。

可以通过以下方式开始对大小为 256x256 的图像进行训练：

  1.  python main.py --img_size 256 --resume_iter 100 --mode train --num_domains 3 --w_hpf 0 \
  2.    --lambda_reg 1 --lambda_sty 1 --lambda_ds 2 --lambda_cyc 1 \
  3.    --train_img_dir data/afhq/train --val_img_dir data/afhq/val \
  4.    --batch_size 4 --sample_every 100 --save_every 500

要在 512x512px 分辨率上进行训练，请运行：

  1.  python main.py --img_size 256 --resume_iter 100 --mode train --num_domains 3 --w_hpf 0 \
  2.    --lambda_reg 1 --lambda_sty 1 --lambda_ds 2 --lambda_cyc 1 \
  3.    --train_img_dir data/afhq/train --val_img_dir data/afhq/val \
  4.    --batch_size 4 --sample_every 100 --save_every 500

本文的源代码：https://github.com/Hramchenko/modulated_stargan

引用

  1.  Yunjey Choi, Youngjung Uh, Jaejun Yoo, Jung-Woo Ha, StarGAN v2: Diverse Image Synthesis for Multiple Domains.
  2.  Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun, Deep Residual Learning for Image Recognition
  3.  Tero Karras, Samuli Laine, Timo Aila, A Style-Based Generator Architecture for Generative Adversarial Networks.
  4.  Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, Timo Aila, Analyzing and Improving the Image Quality of StyleGAN
  5.  https://github.com/lucidrains/stylegan2-pytorch/

本文作者：Vitaliy Hramchenko

标签：
登录后发布评论

“使用Modulated Convolutions修改 StarGAN V2”的评论:

还没有评论
关于作者
...
Deephub
更多文章请关注公众号：Deephub-IMBA
相关阅读

频率派与贝叶斯统计在营销组合建模中的应用比较：隐私优先时代的方法选择

基于Transformer架构的时间序列数据去噪技术研究

英伟达新一代GPU架构（50系列显卡）PyTorch兼容性解决方案

FlashTokenizer: 基于C++的高性能分词引擎，速度可以提升8-15倍

计算加速技术比较分析：GPU、FPGA、ASIC、TPU与NPU的技术特性、应用场景及产业生态

标签噪声下的模型评估：如何准确评估AI模型的真实性能，提高模型性能测量的可信度

VideoMind：Chain-of-LoRA突破时间盲区让AI真正看懂长视频