1.DataLoader 使用多線程加載輸入，設(shè)置num_workers

if args.distributed:

? ? ? ? train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)

? ? else:

? ? ? ? train_sampler = None

? ? train_loader = torch.utils.data.DataLoader(

? ? ? ? train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),

? ? ? ? num_workers=args.workers, pin_memory=True, sampler=train_sampler)

2.加載數(shù)據(jù)輸入到CUDA 設(shè)備時設(shè)置非堵塞 non_blocking=True

if args.gpu is not None:

? ? ? ? ? ? input = input.cuda(args.gpu, non_blocking=True)

? ? ? ? ? ? target = target.cuda(args.gpu, non_blocking=True)

3.使用nvidia DALI 加速load 數(shù)據(jù)

準備pipeline：

pipe = HybridValPipe(batch_size=1280,num_threads=4,device_id=0,

data_dir=testdir,crop=64,local_rank=0,world_size=1,

size=64)

pipe.build()

test_loader = DALIClassificationIterator(pipe,size=int(pipe.epoch_size("Reader") /1))

詳細見：https://docs.nvidia.com/deeplearning/dali/user-guide/docs/api.html