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的作用,它就是突出了边缘,而弱化了别的部分,这样就可以让模型更关注于边缘的部分,对于神经网络这种黑盒模型来说,加强了模型的方向性,可以取得不错的效果。