update comment

ChenJian7578 · Mar 9, 2022 · 1fff15a · 1fff15a
1 parent 44162a7
commit 1fff15a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/train.py b/train.py
@@ -67,10 +67,9 @@
     #   此处使用的是整个模型的权重，因此是在train.py进行加载的，下面的pretrain不影响此处的权值加载。
     #   如果想要让模型从主干的预训练权值开始训练，则设置model_path = ''，下面的pretrain = True，此时仅加载主干。
     #   如果想要让模型从0开始训练，则设置model_path = ''，下面的pretrain = Fasle，Freeze_Train = Fasle，此时从0开始训练，且没有冻结主干的过程。
-    #   一般来讲，从0开始训练效果会很差，因为权值太过随机，特征提取效果不明显。
-    #
-    #   网络一般不从0开始训练，至少会使用主干部分的权值，有些论文提到可以不用预训练，主要原因是他们 数据集较大 且 调参能力优秀。
-    #   如果一定要训练网络的主干部分，可以了解imagenet数据集，首先训练分类模型，分类模型的 主干部分 和该模型通用，基于此进行训练。
+    #   
+    #   一般来讲，网络从0开始的训练效果会很差，因为权值太过随机，特征提取效果不明显，因此非常、非常、非常不建议大家从0开始训练！
+    #   如果一定要从0开始，可以了解imagenet数据集，首先训练分类模型，获得网络的主干部分权值，分类模型的 主干部分 和该模型通用，基于此进行训练。
     #----------------------------------------------------------------------------------------------------------------------------#
     model_path      = 'model_data/yolo_weights.pth'
     #------------------------------------------------------#

diff --git a/utils/callbacks.py b/utils/callbacks.py
@@ -18,8 +18,11 @@ def __init__(self, log_dir, model, input_shape):
 
         os.makedirs(self.log_dir)
         self.writer     = SummaryWriter(self.log_dir)
-        dummy_input     = torch.randn(2, 3, input_shape[0], input_shape[1])
-        self.writer.add_graph(model, dummy_input)
+        try:
+            dummy_input     = torch.randn(2, 3, input_shape[0], input_shape[1])
+            self.writer.add_graph(model, dummy_input)
+        except:
+            pass
 
     def append_loss(self, epoch, loss, val_loss):
         if not os.path.exists(self.log_dir):