From dfa7e047a484cbaf7464d3392351d05832dea43a Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 5 Aug 2019 13:55:42 +0200 Subject: [PATCH] updates --- test.py | 2 +- train.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/test.py b/test.py index 9dbb738e..f98811fe 100644 --- a/test.py +++ b/test.py @@ -48,7 +48,7 @@ def test(cfg, dataset = LoadImagesAndLabels(test_path, img_size, batch_size) dataloader = DataLoader(dataset, batch_size=batch_size, - num_workers=4, + num_workers=os.cpu_count(), pin_memory=True, collate_fn=dataset.collate_fn) diff --git a/train.py b/train.py index 4678e4ba..8234a0ff 100644 --- a/train.py +++ b/train.py @@ -7,7 +7,6 @@ import torch.optim.lr_scheduler as lr_scheduler import test # import test.py to get mAP after each epoch from models import * -from utils.adabound import * from utils.datasets import * from utils.utils import * @@ -344,7 +343,7 @@ if __name__ == '__main__': parser.add_argument('--rect', action='store_true', help='rectangular training') parser.add_argument('--resume', action='store_true', help='resume training flag') parser.add_argument('--transfer', action='store_true', help='transfer learning flag') - parser.add_argument('--num-workers', type=int, default=4, help='number of Pytorch DataLoader workers') + parser.add_argument('--num-workers', type=int, default=os.cpu_count(), help='DataLoader workers') parser.add_argument('--nosave', action='store_true', help='only save final checkpoint') parser.add_argument('--notest', action='store_true', help='only test final epoch') parser.add_argument('--xywh', action='store_true', help='use xywh loss instead of GIoU loss')