Increase training performance

pull/1/head
Ross Wightman 5 years ago
parent 5855b07ae0
commit 72b4d162a2

@ -62,7 +62,7 @@ parser.add_argument('--log-interval', type=int, default=50, metavar='N',
help='how many batches to wait before logging training status') help='how many batches to wait before logging training status')
parser.add_argument('--recovery-interval', type=int, default=1000, metavar='N', parser.add_argument('--recovery-interval', type=int, default=1000, metavar='N',
help='how many batches to wait before writing recovery checkpoint') help='how many batches to wait before writing recovery checkpoint')
parser.add_argument('-j', '--workers', type=int, default=2, metavar='N', parser.add_argument('-j', '--workers', type=int, default=6, metavar='N',
help='how many training processes to use (default: 1)') help='how many training processes to use (default: 1)')
parser.add_argument('--num-gpu', type=int, default=1, parser.add_argument('--num-gpu', type=int, default=1,
help='Number of GPUS to use') help='Number of GPUS to use')
@ -113,7 +113,6 @@ def main():
loader_train = data.DataLoader( loader_train = data.DataLoader(
dataset_train, dataset_train,
batch_size=batch_size, batch_size=batch_size,
pin_memory=True,
shuffle=True, shuffle=True,
num_workers=args.workers num_workers=args.workers
) )
@ -125,7 +124,6 @@ def main():
loader_eval = data.DataLoader( loader_eval = data.DataLoader(
dataset_eval, dataset_eval,
batch_size=4 * args.batch_size, batch_size=4 * args.batch_size,
pin_memory=True,
shuffle=False, shuffle=False,
num_workers=args.workers num_workers=args.workers
) )

Loading…
Cancel
Save