这是训练250epoch左右的成果。
之前的文章里面，我们使用了残差网络的形式实现生成器与辨别器，它理论上可以实现很不错的效果，但有一个很致命的缺点，就是训练太慢，很难见到成果。
这一次，我们实现了一个利用自注意力机制制作的对抗生成网络。自注意力机制是我们在深度学习道路上，除了RNN,CNN以外，不得不了解的一种模块。非常有意思。简而言之，这个模块，相比于之前单纯使用卷积网络的GAN，它更加能注重上下文，举个例子，在生成人物眼睛的时候，它会注意到鼻子👃，头发等其他部位，从而将眼睛放在合适的位置，总之，能更好的学习到整体特征。
以下是生成器的代码。注意attn1与attn2层都是我们的自注意力模块，其他是我们所熟悉的DCGAN中使用过的反卷积。默认生成64*64像素的图片，如果想修改图片大小，请修改image_size，以及，后面层次中的channel大小以及selfattention的参数。
class Generator(nn.Module):def __init__(self, image_size = 64, z_dim = 100, conv_dim =64):super().__init__()repeat_num = int(np.log2(image_size)) - 3mult = 2 ** repeat_numself.l1 = nn.Sequential(spectral_norm(nn.ConvTranspose2d(in_channels = z_dim, out_channels = conv_dim * mult, kernel_size = 4)),nn.LayerNorm([512, 4, 4]),nn.ReLU())curr_dim = conv_dim * multself.l2 = nn.Sequential(spectral_norm(nn.ConvTranspose2d(curr_dim, curr_dim // 2, 4, 2, 1)),nn.LayerNorm([256, 8, 8]),nn.ReLU())curr_dim = curr_dim // 2self.l3 = nn.Sequential(spectral_norm(nn.ConvTranspose2d(curr_dim, curr_dim // 2, 4, 2, 1)),nn.LayerNorm([128, 16, 16]),nn.ReLU())curr_dim = curr_dim // 2self.l4 = nn.Sequential(spectral_norm(nn.ConvTranspose2d(curr_dim, curr_dim // 2, 4, 2, 1)),nn.LayerNorm([64, 32, 32]),nn.ReLU())self.last = nn.Sequential(nn.ConvTranspose2d(64, 3, 4, 2, 1),nn.Tanh())self.attn1 = selfattention(128)self.attn2 = selfattention(64)def forward(self, input):input = input.view(input.size(0), input.size(1), 1, 1)out = self.l1(input)out = self.l2(out)out = self.l3(out)out = self.attn1(out)out = self.l4(out)out = self.attn2(out)out = self.last(out)return out
以下是辨别器的代码，同样，如果上面修改了数据集图片读入大小以及生成器生成的图片大小，不要忘记修改辨别器中的image_size，同时可能还需要修改每个卷积层的核大小以及步长。
class Discriminator(nn.Module):def __init__(self, in_channels = 3, image_size = 256, ndf =64):super().__init__()def conv_2d(in_channels, out_channels, kernel_size, stride = 1, padding = 0):return nn.Sequential(spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)),nn.LeakyReLU(0.1))self.block_1 = conv_2d(in_channels, ndf, 4, 2, 1)current_dim = ndfself.block_2 = conv_2d(current_dim, current_dim * 2, 4, 2, 1)current_dim *= 2self.block_3 = conv_2d(current_dim, current_dim * 2, 4, 2, 1)current_dim *= 2#self.block_5 = conv_2d(current_dim, current_dim * 2, 4, 2, 1)# current_dim *= 2#self.block_6 = conv_2d(current_dim, current_dim * 2, 4, 2, 1)#current_dim *= 2self.attn_layer_1 = selfattention(current_dim)self.block_4 = conv_2d(current_dim, current_dim * 2, 4, 2, 1)current_dim *= 2self.attn_layer_2 = selfattention(current_dim)self.last_layer = nn.Sequential(nn.Conv2d(current_dim, 1, 4, stride= 1),)def forward(self, input):all_layers = [self.block_1, self.block_2, self.block_3, self.attn_layer_1,self.block_4, self.attn_layer_2,self.last_layer]out = reduce(lambda x, layer: layer(x), all_layers, input) #套娃 clock3(block2(block1(x)))......返回结果return out
以下贴出自注意力机制代码，有疑问的同学可以参考之前的博客:
class selfattention(nn.Module):def __init__(self, in_channels):super().__init__()self.in_channels = in_channelsself.query = nn.Conv2d(in_channels, in_channels // 8, kernel_size = 1, stride = 1)self.key = nn.Conv2d(in_channels, in_channels // 8, kernel_size = 1, stride = 1)self.value = nn.Conv2d(in_channels, in_channels, kernel_size = 1, stride = 1)self.gamma = nn.Parameter(torch.zeros(1)) #gamma为一个衰减参数，由torch.zero生成，nn.Parameter的作用是将其转化成为可以训练的参数.self.softmax = nn.Softmax(dim = -1)def forward(self, input):batch_size, channels, height, width = input.shape# input: B, C, H, W -> q: B, H * W, C // 8q = self.query(input).view(batch_size, -1, height * width).permute(0, 2, 1)#input: B, C, H, W -> k: B, C // 8, H * Wk = self.key(input).view(batch_size, -1, height * width)#input: B, C, H, W -> v: B, C, H * Wv = self.value(input).view(batch_size, -1, height * width)#q: B, H * W, C // 8 x k: B, C // 8, H * W -> attn_matrix: B, H * W, H * Wattn_matrix = torch.bmm(q, k) #torch.bmm进行tensor矩阵乘法,q与k相乘得到的值为attn_matrix.attn_matrix = self.softmax(attn_matrix)#经过一个softmax进行缩放权重大小.out = torch.bmm(v, attn_matrix.permute(0, 2, 1)) #tensor.permute将矩阵的指定维进行换位.这里将1于2进行换位。out = out.view(*input.shape)return self.gamma * out + input
训练代码以及损失函数等全部源代码，请上我的GitHub获取，之后会上传预训练模型，您也可以自己进行训练~~~///(^v^)\\\~~~
数据集来自kaggle，可自己搜索："anime"然后下载。
Github链接：/rabbitdeng/anime-sagan-pytorch/tree/main
