From a2ad00d6fccbaa3ce3026abef8136fd9f1a607e1 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 5 Mar 2019 17:10:34 +0100 Subject: [PATCH] updates --- train.py | 7 ++++--- utils/utils.py | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/train.py b/train.py index 099730f3..e12a0937 100644 --- a/train.py +++ b/train.py @@ -85,8 +85,9 @@ def train( # Set scheduler # scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[54, 61], gamma=0.1) - model_info(model) t0 = time.time() + model_info(model) + n_burnin = min(dataloader.nB / 5, 1000) # number of burn-in batches for epoch in range(epochs): epoch += start_epoch @@ -118,8 +119,8 @@ def train( continue # SGD burn-in - if (epoch == 0) & (i <= 1000): - lr = lr0 * (i / 1000) ** 4 + if (epoch == 0) & (i <= n_burnin): + lr = lr0 * (i / n_burnin) ** 4 for g in optimizer.param_groups: g['lr'] = lr diff --git a/utils/utils.py b/utils/utils.py index a96467bd..ebbea4d3 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -379,7 +379,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4): if prediction.is_cuda: unique_labels = unique_labels.cuda(prediction.device) - nms_style = 'MERGE' # 'OR' (default), 'AND', 'MERGE' (experimental) + nms_style = 'OR' # 'OR' (default), 'AND', 'MERGE' (experimental) for c in unique_labels: # Get the detections with class c dc = detections[detections[:, -1] == c]