From 3bea4da6040189f794f2c6cb41083151f444d61c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 4 Mar 2019 17:34:53 +0100 Subject: [PATCH] updates --- models.py | 6 +----- train.py | 4 ++-- utils/torch_utils.py | 2 +- utils/utils.py | 8 -------- 4 files changed, 4 insertions(+), 16 deletions(-) diff --git a/models.py b/models.py index 8f7f79d6..ded1ae71 100755 --- a/models.py +++ b/models.py @@ -150,10 +150,7 @@ class YOLOLayer(nn.Module): p_conf = p[..., 4] # Conf p_cls = p[..., 5:] # Class - if p.is_cuda: - txy, twh, mask, tcls = build_targets(targets, self.anchor_vec.cuda(), self.nA, self.nC, nG) - else: - txy, twh, mask, tcls = build_targets(targets, self.anchor_vec, self.nA, self.nC, nG) + txy, twh, mask, tcls = build_targets(targets, self.anchor_vec, self.nA, self.nC, nG) tcls = tcls[mask] if p.is_cuda: @@ -164,7 +161,6 @@ class YOLOLayer(nn.Module): nM = mask.sum().float() # number of anchors (assigned to targets) k = 1 # nM / bs if nM > 0: - print(xy.shape, txy.shape, mask.shape) lxy = k * MSELoss(xy[mask], txy[mask]) lwh = k * MSELoss(wh[mask], twh[mask]) diff --git a/train.py b/train.py index 59ac48bc..57b5d44b 100644 --- a/train.py +++ b/train.py @@ -48,8 +48,8 @@ def train( # Load weights to resume from model.load_state_dict(checkpoint['model']) - # if torch.cuda.device_count() > 1: - # model = nn.DataParallel(model) + if torch.cuda.device_count() > 1: + model = nn.DataParallel(model) model.to(device).train() # Transfer learning (train only YOLO layers) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index dc7a4e69..504aeba2 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -17,7 +17,7 @@ def select_device(force_cpu=False): if torch.cuda.device_count() > 1: print('Found %g GPUs' % torch.cuda.device_count()) - print('WARNING Using GPU0 Only: https://github.com/ultralytics/yolov3/issues/21') + print('WARNING Multi-GPU Issue: https://github.com/ultralytics/yolov3/issues/21') # torch.cuda.set_device(0) # OPTIONAL: Set your GPU if multiple available # # print('Using ', torch.cuda.device_count(), ' GPUs') diff --git a/utils/utils.py b/utils/utils.py index c214fd2d..a96467bd 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -242,12 +242,6 @@ def build_targets(target, anchor_vec, nA, nC, nG): tconf = torch.ByteTensor(nB, nA, nG, nG).fill_(0) tcls = torch.ByteTensor(nB, nA, nG, nG, nC).fill_(0) # nC = number of classes - if anchor_vec.is_cuda: - txy = txy.cuda() - twh = twh.cuda() - tconf = tconf.cuda() - tcls = tcls.cuda() - for b in range(nB): t = target[b] nTb = len(t) # number of targets @@ -263,8 +257,6 @@ def build_targets(target, anchor_vec, nA, nC, nG): box1 = gwh box2 = anchor_vec.unsqueeze(1) - print(box1.device, box2.device) - inter_area = torch.min(box1, box2).prod(2) iou = inter_area / (box1.prod(1) + box2.prod(2) - inter_area + 1e-16)