pytorch 训练加速Tips

1.DataLoader 使用多线程加载输入，设置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.加载数据输入到CUDA 设备时设置非堵塞 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 数据

准备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