【pytorch】梯度爆炸/消失解决办法

every blog every motto: You can do more than you think.

0. 前言

记录一种解决梯度爆炸/消失的解决办法，后续出现类似情况再进行增补

1. 正文

1.1 介绍

1.1.1 背景

训练过程中，出现损失为nan了，知道情况不妙了，如下

1.1.2 用tensorboard记录下梯度情况

说明： 数据比较大，所以 在每次迭代过程中记录
注意： 每次迭代是耗时的，因为会将损失从GPU传递到CPU，然后再进行记录，我们通过nvidia-smi查看GPU使用情况，发现GPU利用率不高，遂将改为每间隔20次进行保存。

 writer_iter_dir = r'runs/v1/iter'  # 每次迭代保存一次
 if not os.path.join(writer_iter_dir): os.makedirs(writer_iter_dir)
 self.writer_iter = SummaryWriter(writer_iter_dir)
 
for epoch in range(1,50):
	for iter,data in enumerate(dataloader):
        if i % 20 == 0:
           self.writer_iter.add_scalar('Loss/Train_iter', batch_all_loss.item(), i)  # 保存每个batch的损失
           
           for name, param in self.net.named_parameters():
               self.writer_iter.add_histogram('batch_' + name + '_param', param, i)  # 记录参数
               self.writer_iter.add_histogram('batch_' + name + '_grad', param.grad, i)  # 记录梯度

1.1.3 查看损失和梯度

损失如下：

如下图所示，
我们查看的是每次迭代的记录（我们也进行了每个epoch的保存，代码和上面类似，不赘述）
我们看到当前的卷积块：stage_1_decoding.double_conv.0，其权重 的梯度 在第1次（下图显示的迭代次数为0）为2.05e4，后续迭代也是类似

我们再看下权重，如下图所示
权重也是比较大，第160次迭代为2.93e3，且为均匀分布

1.2 解决办法

梯度爆炸可能和我们使用Pytorch中卷积的层的权重有关，其中，默认权重为均匀分布，如下图所示
pytorch，卷积层链接

我们对其进行修改为正态分布，
说明： 模型要先实例化

 # 权重初始化
 for m in self.net.modules():
     if isinstance(m, nn.Conv2d):
         init.xavier_normal(m.weight.data)

修改后的损失，

修改后的权重如下

梯度如下，

参考文献

[1] https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html?highlight=conv2d#torch.nn.Conv2d
[2] https://tensorflow.google.cn/versions/r2.0/api_docs/python/tf/keras/layers/Conv2D
[3] https://blog.csdn.net/u010368556/article/details/81148267
[4] https://blog.csdn.net/qq_25737169/article/details/78847691
[5] https://www.cnblogs.com/zhangxiann/p/13617142.html
[6] https://www.cnblogs.com/zhangxiann/p/13617142.html
[7] https://blog.csdn.net/u010099080/article/details/77426577
[8] https://zhuanlan.zhihu.com/p/56253634
[9] https://www.cnblogs.com/my-love-is-python/p/11725664.html