PraNet训练文件train.py详解——shaoshuai

1、train.py从main函数开始，首先规定了参数。

以上参数大多都是见名知意的，下面对某些重要参数做解释。

• --lr:学习率。学习率的大小代表梯度下降的快慢，在使用Adam的优化函数时，lr的默认初始值为e-4。
• --batchsize。每次训练时一个批的数据中包含多少张图片。
• --clip。梯度裁剪。这个参数用在对梯度下降的梯度进行裁剪，这时为了防止梯度爆炸。将梯度限制在一定范围内。

2、输入增强。

为了增强模型的鲁棒性，增加可输入的图片，这里对图片进行增强处理。

3、模型训练

接下来进行的就是模型的训练，训练模型时会返回4个预测值。

 这分别对应着论文中标红的地方。

 

以下是在实际运行过程中的四张图。

 

从左上到右下依次为lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2。

4、损失函数计算。

损失函数的计算采用的是二至交叉熵加iou交并比的方法，与传统方法不同的是，PraNet在这里加入边缘权重的概念，加强了对边缘的关注。

 

我们可以看到weit的作用，它就是突出了边缘，而弱化了别的部分，这样就可以让模型更关注于边缘的部分，对于神经网络这种黑盒模型来说，加强了模型的方向性，可以取得不错的效果。