PyTorch代码优化技巧

1 直接在GPUs上构建张量

很多人都是这样在GPUs上创建张量的

t = tensor.rand(2,2).cuda()

然而，这首先创建CPU张量，然后将其转移到GPU……这真的很慢。相反，直接在想要的设备上创建张量。

t = tensor.rand(2,2, device=torch.device('cuda:0'))

2 使用DistributedDataParallel不要使用DataParallel

PyTorch有两个主要的模式用于在多 GPUs训练。
第一种是DataParallel，它将一批数据分割到多个GPUs上。但这也意味着模型必须复制到每个GPU上，一旦在GPU 0上计算出梯度，它们必须同步到其他GPU。
第二种是DistributedDataParallel在每个GPU(在它自己的进程中)上创建模型副本，并且只让数据的一部分对该GPU可用。这就像是让N个独立的模型进行训练，除了一旦每个模型都计算出梯度，它们就会在模型之间同步梯度……这意味着我们在每批处理中只在GPUs之间传输一次数据。

3 DataLoaders中workers的设置

推荐
num_worker = 4 * num_GPU

4 DataLoader中启用pinned_memory时，避免调用torch.cuda.empty_cache()

当你在一个DataLoader中启用pinned_memory时，它“自动将获取的数据张量放在pinned memory中，并使数据更快地传输到CUDA-enabled的gpu”

5 避免CPU和GPU之间的数据传输

避免使用
.item()
.numpy()
因为将数据从GPU传输到CPU，从而极大地降低了性能。

如果你试图清除附加的计算图，建议使用.detach()。这不会将内存转移到GPU，它会删除任何附加到该变量的计算图。