4.2.7. 模型权重初始化#

模型权重初始化会对我们模型训练造成很大影响。初始权重在一定程度上决定了模型的收敛情况,合理的权重可以加快梯度下降收敛的速度。若权重过大,将造成梯度爆炸,而权重过小,则易造成信息丢失。我们对生成器中的ConvTranspose2d、BatchNorm2d和判别器中的Conv2d、BatchNorm2d进行参数初始化,使用PyTorch中的nn.init.normal_方法和 nn.init.constant_方法。前者将所选参数按照设定值,进行符合正态分布的权重初始化;后者将所选参数按照设定的常数进行初始化。如下示代码,对于卷积层,其权重按均值为0,标准差为0.02的正态分布进行初始化;对于BatchNorm层,其权重被初始化为均值为1.0,标准差为0.02的正态分布下的随机值,并设置偏置为常数0。

生成器网络和判别器网络初始化后,需分配到设备device中去。device是我们先前设置参数时得到的表示设备的变量,如果我们使用的计算机支持CUDA,那么我们可以把神经网络模型和张量都通过to(device)方法放到GPU里训练,这将大大提高训练速度。

chapter_4_2_7_01.py#
 1from torch import nn
 2
 3from chapter_4_2_3_01 import device
 4from chapter_4_2_5_01 import Generator
 5from chapter_4_2_6_01 import Discriminator
 6
 7
 8def weights_init(m):
 9    """
10    模型权重初始化
11    Args:
12        m:
13
14    Returns:
15
16    """
17    classname = m.__class__.__name__
18    if classname.find('Conv') != -1:
19        nn.init.normal_(m.weight.data, 0.0, 0.02)
20    elif classname.find('BatchNorm') != -1:
21        nn.init.normal_(m.weight.data, 1.0, 0.02)
22        nn.init.constant_(m.bias.data, 0)
23
24
25netG = Generator().to(device)
26netD = Discriminator().to(device)
27netG.apply(weights_init)
28netD.apply(weights_init)