diff --git a/train.py b/train.py index 58403cf..a940ec7 100644 --- 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 index 4a7f597..039daaa 100644 --- 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):