updates
This commit is contained in:
parent
48af6d136f
commit
1c0d408fbf
2
test.py
2
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=os.cpu_count(),
|
||||
num_workers=min(os.cpu_count(), batch_size),
|
||||
pin_memory=True,
|
||||
collate_fn=dataset.collate_fn)
|
||||
|
||||
|
|
3
train.py
3
train.py
|
@ -172,7 +172,7 @@ def train(cfg,
|
|||
# Dataloader
|
||||
dataloader = torch.utils.data.DataLoader(dataset,
|
||||
batch_size=batch_size,
|
||||
num_workers=opt.num_workers,
|
||||
num_workers=min(os.cpu_count(), batch_size),
|
||||
shuffle=not opt.rect, # Shuffle=True unless rectangular training is used
|
||||
pin_memory=True,
|
||||
collate_fn=dataset.collate_fn)
|
||||
|
@ -341,7 +341,6 @@ 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=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')
|
||||
|
|
Loading…
Reference in New Issue