BigGAN代码解读(gpt4.0的帮助)——谱正则化部分

作者个人记录学习

BigGAN中使用谱归一化对训练过程进行优化，在github中的代码中，使用了自己编写的谱归一化对卷积层、线性层以及Embedding层进行重写，下面对这部分重写进行解读。
代码网址：点赞最高的BigGAN代码

在这片代码中说自己是officially unoffcail，不是很理解

谱归一化代码如下（英文为原本注释，中文为作者后续理解注释）：

# Spectral normalization base class 
class SN(object):
  def __init__(self, num_svs, num_itrs, num_outputs, transpose=False, eps=1e-12):
    # Number of power iterations per step
    # 幂迭代的迭代次数,幂迭代用于估计特征值
    self.num_itrs = num_itrs
    # Number of singular values
    # 想要估计的特征值数量
    self.num_svs = num_svs
    # Transposed?
    self.transpose = transpose
    # Epsilon value for avoiding divide-by-0
    self.eps = eps
    # Register a singular vector for each sv
    for i in range(self.num_svs):
      # 缓冲区
      # 特征向量
      self.register_buffer('u%d' % i, torch.randn(1, num_outputs))
      # 特征值
      self.register_buffer('sv%d' % i, torch.ones(1))
  
  # Singular vectors (u side)
  # 保存特征向量
  @property
  def u(self):
    return [getattr(self, 'u%d' % i) for i in range(self.num_svs)]

  # Singular values; 
  # note that these buffers are just for logging and are not used in training. 
  # 保存特征值
  @property
  def sv(self):
   return [getattr(self, 'sv%d' % i) for i in range(self.num_svs)]
   
  # Compute the spectrally-normalized weight
  def W_(self):
    W_mat = self.weight.view(self.weight.size(0), -1)
    if self.transpose:
      W_mat = W_mat.t()
    # Apply num_itrs power iterations
    for _ in range(self.num_itrs):
      svs, us, vs = power_iteration(W_mat, self.u, update=self.training, eps=self.eps) 
    # Update the svs
    if self.training:
      with torch.no_grad(): # Make sure to do this in a no_grad() context or you'll get memory leaks!
        for i, sv in enumerate(svs):
          self.sv[i][:] = sv  
    # 返回谱归一化的权重矩阵,即原权重矩阵除以最大特征值的近似值   
    return self.weight / svs[0]

这部分代码用于谱正则化，主要的函数W_首先计算权重矩阵的特征值的近似值，计算方法为幂迭代，之后使用原本的权重矩阵除以最大的特征值，并返回调整之后的权重矩阵。

重写的卷积层、线性层、Embedding层代码：

# 2D Conv layer with spectral norm
class SNConv2d(nn.Conv2d, SN):
  # 继承自nn.Conv2d, SN
  def __init__(self, in_channels, out_channels, kernel_size, stride=1,
             padding=0, dilation=1, groups=1, bias=True, 
             num_svs=1, num_itrs=1, eps=1e-12):
    nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, stride, 
                     padding, dilation, groups, bias)
    SN.__init__(self, num_svs, num_itrs, out_channels, eps=eps)    
  def forward(self, x):
    # self.W_()来源于谱归一化后的权重
    return F.conv2d(x, self.W_(), self.bias, self.stride, 
                    self.padding, self.dilation, self.groups)

# 线性层和embedding层和卷积层一样
# Linear layer with spectral norm
class SNLinear(nn.Linear, SN):
  def __init__(self, in_features, out_features, bias=True,
               num_svs=1, num_itrs=1, eps=1e-12):
    nn.Linear.__init__(self, in_features, out_features, bias)
    SN.__init__(self, num_svs, num_itrs, out_features, eps=eps)
  def forward(self, x):
    return F.linear(x, self.W_(), self.bias)


# Embedding layer with spectral norm
# We use num_embeddings as the dim instead of embedding_dim here
# for convenience sake
class SNEmbedding(nn.Embedding, SN):
  def __init__(self, num_embeddings, embedding_dim, padding_idx=None, 
               max_norm=None, norm_type=2, scale_grad_by_freq=False,
               sparse=False, _weight=None,
               num_svs=1, num_itrs=1, eps=1e-12):
    nn.Embedding.__init__(self, num_embeddings, embedding_dim, padding_idx,
                          max_norm, norm_type, scale_grad_by_freq, 
                          sparse, _weight)
    SN.__init__(self, num_svs, num_itrs, num_embeddings, eps=eps)
  def forward(self, x):
    return F.embedding(x, self.W_())

以卷积层为例，新的卷积层继承了谱归一化类与原本的nn.Conv2d类，就是将nn.Conv2d中的self.weight权重矩阵经过SN的处理，返回的新权重矩阵，在forward函数中，同其他nn.Conv2d原本的参数，例如self.bias, self.stride，一同输入给F.conv2d，形成一个经过谱归一化的二维卷积，后面两个重写也为同理。

感谢gpt4.0帮我看懂。