Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2019-03-05 17:10:34 +01:00
parent 2c2d7bc63b
commit a2ad00d6fc
2 changed files with 5 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
train.py
View File

@ -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

2
utils/utils.py
View File

@ -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]