This commit is contained in:
Glenn Jocher 2019-12-08 17:57:23 -08:00
parent e35397ee41
commit 61c3cb9ecf
1 changed files with 2 additions and 2 deletions

View File

@ -199,7 +199,6 @@ def train():
# Dataloader
batch_size = min(batch_size, len(dataset))
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers
print('Using %g dataloader workers' % nw)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
num_workers=nw,
@ -224,11 +223,12 @@ def train():
model.arc = opt.arc # attach yolo architecture
model.hyp = hyp # attach hyperparameters to model
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weights
torch_utils.model_info(model, report='summary') # 'full' or 'summary'
maps = np.zeros(nc) # mAP per class
# torch.autograd.set_detect_anomaly(True)
results = (0, 0, 0, 0, 0, 0, 0) # 'P', 'R', 'mAP', 'F1', 'val GIoU', 'val Objectness', 'val Classification'
t0 = time.time()
torch_utils.model_info(model, report='summary') # 'full' or 'summary'
print('Using %g dataloader workers' % nw)
print('Starting %s for %g epochs...' % ('prebias' if opt.prebias else 'training', epochs))
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------
model.train()