PyTorch 常用代码段整理合集

多卡同步BN

当使用torch.nn.DataParallel将代码运行在多张GPU卡上时，PyTorch的BN层默认操作是各卡上数据独立地计算均值和标准差，同步BN使用所有卡上的数据一起计算BN层的均值和标准差，缓解了当批量大小（batch size）比较小时对均值和标准差估计不准的情况，是在目标检测等任务中一个有效的提升性能的技巧。

Github:Synchronized-BatchNorm-PyTorch

Pytorch官方现已支持
备注：虽然官方已经支持了，但是文档中有说，目前只支持DistributedDataParallel，我自己在单机多卡上测试时一直报错，说未初始化，所以还是尽量有上面Github那个吧，亲测可用(＾Ｕ＾)ノ~ＹＯ

sync_bn = torch.nn.SyncBatchNorm(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)

固定随机种子

torch.manual_seed(0)
torch.cuda.manual_seed_all(0)

计算模型参数量

param_count = sum(param.numel() for param in model.parameters())

现在Github上有开源的工具可以直接调用 计算模型的计算量和参数量，传送门：flops-counter.pytorch

import torchvision.models as models
import torch
from ptflops import get_model_complexity_info

with torch.cuda.device(0):
  net = models.densenet161()
  flops, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True, print_per_layer_stat=True)
  print('Flops:  ' + flops)
  print('Params: ' + params)

提升Pytorch运行效率

参考自：pytorch torch.backends.cudnn设置作用
cuDNN使用非确定性算法，并且可以使用torch.backends.cudnn.enabled = False来进行禁用。

如果设置为torch.backends.cudnn.enabled =True，说明设置为使用使用非确定性算法。

然后再设置：torch.backends.cudnn.benchmark = true，那么cuDNN使用的非确定性算法就会自动寻找最适合当前配置的高效算法，来达到优化运行效率的问题。

一般来讲，应该遵循以下准则：

• 如果网络的输入数据维度或类型上变化不大，设置 torch.backends.cudnn.benchmark = true 可以增加运行效率；
• 如果网络的输入数据在每次 iteration 都变化的话，会导致 cnDNN 每次都会去寻找一遍最优配置，这样反而会降低运行效率。

所以我们经常看见在代码开始出两者同时设置：

torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True

指定程序运行在特定GPU卡上

在命令行指定环境变量

CUDA_VISIBLE_DEVICES=0,1 python train.py

或在代码中指定

os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

保证模型的可重复性

参考：pytorch如何能够保证模型的可重复性
GPU上两次运行的结果会不一致，而CPU两次运行的结果相同，解决方法在主函数中加上一句

torch.backends.cudnn.deterministic = True