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帮我看懂。