From e1c407dab1841008daf4cd50d8c61d78a76035f3 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 5 Aug 2019 17:25:50 +0200 Subject: [PATCH] updates --- train.py | 1 + utils/utils.py | 5 ++--- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/train.py b/train.py index d24a027d..f8f04e58 100644 --- a/train.py +++ b/train.py @@ -178,6 +178,7 @@ def train(cfg, world_size=1, # number of nodes for distributed training rank=0) # distributed training node rank model = torch.nn.parallel.DistributedDataParallel(model) + model.yolo_layers = model.module.yolo_layers # move yolo layer indices to top level # Dataset dataset = LoadImagesAndLabels(train_path, diff --git a/utils/utils.py b/utils/utils.py index 65dbf917..91d75e2c 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -334,13 +334,12 @@ def compute_loss(p, targets, model, giou_loss=True): # predictions, targets, mo def build_targets(model, targets): # targets = [image, class, x, y, w, h] iou_thres = model.hyp['iou_t'] # hyperparameter - if type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel): - model = model.module + multi_gpu = type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel) nt = len(targets) txy, twh, tcls, tbox, indices, anchor_vec = [], [], [], [], [], [] for i in model.yolo_layers: - layer = model.module_list[i] + layer = model.module.module_list[i] if multi_gpu else model.module_list[i] # iou of targets-anchors t, a = targets, []