如何修改PyTorch预训练GAN输入层 跳过浅层向深层输编码

直接按以下步骤操作即可：

• 第一步：拆分生成器网络结构
  先把加载到的预训练生成器实例打印出来，逐层梳理结构，定位到你想要传入编码的深层全连接层位置。这类预训练GAN的生成器一般是堆叠式结构：前段是小维度全连接层，负责把低维初始噪声映射到高维特征空间，后段是大尺寸全连接+上采样/转置卷积块，负责把高维特征转成最终输出图像。
  定位到目标层后，直接把网络拆成两段，丢弃你要跳过的前序层段，保留从目标层到最终输出层的所有结构即可。如果模型是nn.Sequential格式可以直接切片，非Sequential的带自定义命名层的模型，手动把需要的层按顺序组装成新的推理网络就行，参考代码：

  import torch
  from pretrained_gans import load_pretrained_gan
  
  # 加载目标预训练模型
  gan_model = load_pretrained_gan(model_name="BigGAN-deep-256")
  generator = gan_model.generator
  # 打印结构定位目标层位置
  print(generator)
  
  # 示例1：Sequential结构直接切片，假设前3层是要跳过的初始全连接段
  deep_generator = generator[3:]
  
  # 示例2：非Sequential结构手动组装，替换成你模型里对应的层名即可
  # deep_generator = torch.nn.Sequential(
  #     generator.deep_fc,
  #     generator.up_block_1,
  #     generator.up_block_2,
  #     generator.final_conv
  # )
  

• 第二步：确认目标层输入维度
  你自定义传入的编码必须和目标层要求的输入维度完全匹配，最稳妥的确认方式是给目标层注册前向钩子，跑一次原始模型的前向传播，拿到真实的输入张量维度：

  input_shape = []
  def get_input_shape(module, inputs, output):
      input_shape.append(inputs[0].shape)
  
  # 给目标层注册钩子
  target_layer = generator[3] # 替换成你定位到的目标层
  hook = target_layer.register_forward_hook(get_input_shape)
  # 跑一次原始前向
  dummy_z = torch.randn(1, 128) # 替换成模型要求的原始噪声维度
  with torch.no_grad():
      generator(dummy_z)
  hook.remove()
  deep_input_shape = input_shape[0]
  

• 第三步：传入自定义编码直接推理
  构造符合维度要求的编码张量，直接传入切好的深层生成器做前向计算即可：

  # 构造自定义编码，可随机初始化、也可根据目标图像特征优化得到
  custom_code = torch.randn(deep_input_shape)
  deep_generator.eval()
  with torch.no_grad():
      output_img = deep_generator(custom_code)
  

注意事项 #

• 如果模型带残差连接、跨层跳连（比如StyleGAN的风格注入跳连、U-Net结构跳连），不要直接硬切层，需要修改forward逻辑，把跳连路径上的输入也对应替换成你自定义的特征，避免跳连断裂导致输出错乱。
• 如果你的目的是做图像反演匹配目标图，可以把自定义编码设为可训练参数，用生成图和目标图的MSE损失、感知损失做反向传播优化，比优化原始浅层噪声的灵活度更高，适当加L2正则可以减少生成伪影。
• 推理前记得把深层网络设为eval()模式，避免BatchNorm、Dropout层处于训练状态带来的输出扰动。

内容的提问来源于stack exchange，提问作者Mitchell Slapik