Update train.py

Author: bubbliiiing

train.py

@@ -61,8 +61,9 @@
     #   
     #   当model_path = ''的时候不加载整个模型的权值。
     #
-    #   此处使用的是整个模型的权重，因此是在train.py进行加载的。
-    #   如果想要让模型从0开始训练，则设置model_path = ''，下面的Freeze_Train = Fasle，此时从0开始训练，且没有冻结主干的过程。
+    #   此处使用的是整个模型的权重，因此是在train.py进行加载的，下面的pretrain不影响此处的权值加载。
+    #   如果想要让模型从主干的预训练权值开始训练，则设置model_path = ''，下面的pretrain = True，此时仅加载主干。
+    #   如果想要让模型从0开始训练，则设置model_path = ''，下面的pretrain = Fasle，Freeze_Train = Fasle，此时从0开始训练，且没有冻结主干的过程。
     #   一般来讲，从0开始训练效果会很差，因为权值太过随机，特征提取效果不明显。
     #
     #   网络一般不从0开始训练，至少会使用主干部分的权值，有些论文提到可以不用预训练，主要原因是他们 数据集较大 且 调参能力优秀。
@@ -73,6 +74,13 @@
     #   输入的shape大小，一定要是32的倍数
     #------------------------------------------------------#
     input_shape     = [416, 416]
+    #----------------------------------------------------------------------------------------------------------------------------#
+    #   是否使用主干网络的预训练权重，此处使用的是主干的权重，因此是在模型构建的时候进行加载的。
+    #   如果设置了model_path，则主干的权值无需加载，pretrained的值无意义。
+    #   如果不设置model_path，pretrained = True，此时仅加载主干开始训练。
+    #   如果不设置model_path，pretrained = False，Freeze_Train = Fasle，此时从0开始训练，且没有冻结主干的过程。
+    #----------------------------------------------------------------------------------------------------------------------------#
+    pretrained      = True
     #------------------------------------------------------#
     #   Yolov4的tricks应用
     #   mosaic 马赛克数据增强 True or False 
@@ -131,8 +139,9 @@
     #------------------------------------------------------#
     #   创建yolo模型
     #------------------------------------------------------#
-    model = YoloBody(anchors_mask, num_classes)
-    weights_init(model)
+    model = YoloBody(anchors_mask, num_classes, pretrained = pretrained)
+    if not pretrained:
+        weights_init(model)
     if model_path != '':
         #------------------------------------------------------#
         #   权值文件请看README，百度网盘下载