PyTorch 22. PyTorch常用代码段合集

参考： https://zhuanlan.zhihu.com/p/104019160

导入包和版本查询

import torch
import torch.nn as nn
import torchvision
print(torch.__version__)
print(torch.version.cuda)#cuda版本查询
print(torch.backends.cudnn.version())#cudnn版本查询
print(torch.cuda.get_device_name(0))#设备名

可复现性

在硬件设备（CPU,GPU）不同时，完全的可复现性无法保证，即使随机种子相同。但是，在同一个设备上，应该保证可复现性，具体做法是，在程序开始的时候固定torch的随机种子，同时也把numpy的随机种子固定。

np.random.seed(0)
torch.manual_seed(0) #在CPU设置种子用于生成随机数，以使得结果是确定的
torch.cuda.manual_seed_all(0)#为所有的GPU设置种子，以使得结果是确定的

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

显卡设置

如果只需要一张显卡：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

如果需要指定多张显卡，比如0,1显卡

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

也可以在命令行运行代码时设置显卡：

CUDA_VISIBLE_DEVICES=0,1 python train.py

清除显存

torch.cuda.empty_cache()

也可以使用在命令行重置GPU的指令

nvidia-smi --gpu-reset -i [gpu_id]

张量处理

张量的数据类型

PyTorch有9种CPU张量类型和9种GPU张量类型

张量基本信息

tensor = torch.randn(3,4,5)
print(tensor.type()) # 数据类型
print(tensor.size()) # 张量的shape， 是个元组
print(tensor.dim()) # 维度的数量

命名张量

张量命名是一个非常有用的方法，这样可以方便地使用维度的名字来做索引或其他操作，大大提高了可读性、易用性，防止出错。

# 在PyTorch 1.3之前，需要使用注释
# Tensor [N,C,H,W]
images = torch.randn(32,3,56,56)
images.sum(dim=1)
images.select(dim=1, index=0)

# PyTorch 1.3之后
NCHW = ['N', 'C', 'H', 'W']
images = torch.randn(32, 3, 56, 56, names=NCHW)
images.sum('C')
images.select('C', index=0)
# 也可以这么设置
tensor = torch.randn(3,4,1,2,names=('C','N','H','W'))
# 使用align_to可以对维度方便地排序
tensor = tensor.align_to('N','C','H','W')

数据类型转换

# 设置默认类型，pytorch中的floatTensor远远快于DoubleTensor
torch.set_default_tensor_type(torch.FloatTensor)

# 类型转换
tensor = tensor.cuda()
tensor = tensor.cpu()
tensor = tensor.float()
tensor = tensor.long()

torch.Tensor与np.ndarray转换

除了CharTensor，其他所有CPU上的张量都支持转换为numpy格式然后再转换回来。

ndarray = tensor.cpu().numpy()
tensor = torch.from_numpy(ndarray).float()
tensor = torch.from_numpy(ndarray.copy()).float() # if ndarray has negative stride

Torch.tensor与PIL.Image转换

# pytorch中的张量默认采用[N,C,H,W]的顺序，并且数据范围在[0,1]，需要进行转置和规范化
image = PIL.Image.fromarray(torch.clamp(tensor*255, min=0, max=255).byte().permute(1,2,0).cpu().numpy())
image = torchvision.transforms.functional.to_pil_image(tensor)

# PIL.Image -> torch.Tensor
path = r'./figure.jpg'
tensor = torch.from_numpy(np.asarray(PIL.Image.open(path))).permute(2,0,1).float() / 255
tensor = torchvision.transforms.functional.to_tensor(PIL.Image.open(path))

从只包含一个元素的张量中提取值

value = torch.rand(1).item()

张量形变

# 在将卷积层输入全连接层的情况下通常需要对张量做形变处理
# 相比torch.view, torch.reshape可以自动处理输入张量不连续的情况
tensor = torch.rand(2,3,4)
shape=(6, 4)
tensor = torch.reshape(tensor, shape)

打乱顺序

tensor = tensor[torch.randperm(tensor.size(0))] #打乱第一个维度

复制张量

# Operation                 |  New/Shared memory | Still in computation graph |
tensor.clone()            # |        New         |          Yes               |
tensor.detach()           # |      Shared        |          No                |
tensor.detach.clone()()   # |        New         |          No                |

得到非零元素

torch.nonzero(tensor) # 得到非零元素的index
torch.nonzero(tensor==0) # 得到为零元素的index
torch.nonzero(tensor).size(0) # 非零元素的个数
torch.nonzero(tensor == 0).size(0) # 零元素的个数

判断两个张量相等

torch.allclose(tensor1, tensor2) #float tensor
torch.equal(tensor1, tensor2) # int tensor

张量扩展

# 将64*512的张量扩展为64*512*7*7
tensor = torch.rand(64, 512)
torch.reshape(tensor, (64, 512, 1, 1)).expand(64, 512, 7, 7)

多卡同步BN（Batch normalization）

当使用torch.nn.DataParallel将代码运行在多张GPU卡上时，PyTorch的BN层默认操作是各卡上数据独立地计算均值和标准差，同步BN使用所有卡上的数据一起计算BN层的均值和标准差，缓解了bs比较小时对均值和标准差估计不准的情况。

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

将已有网络的所有BN层改为同步BN层

def convertBNtoSyncBN(module, process_group=None):
	# 递归地将网络中所有的BN替换成SyncBN层
	if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
		sync_bn = torch.nn.SyncBatchNorm(module.num_features, module.eps, module.momentum, 
                                         module.affine, module.track_running_stats, process_group)
        sync_bn.running_mean = module.running_mean
        sync_bn.running_var = module.running_var
        if module.affine:
        	sync_bn.weight = module.weight.clone().detach()
        	sync_bn.bias = module.bias.clone().detach()
        return sync_bn
     else:
     	for name, child_module in module.named_children():
     		setattr(module, name) = convert_syncbn_model(child_module, process_group = process_group)
     	return module

affine定义了BN层的参数$\gamma$和$\beta$是否是可学习的（不可学习默认是常数1和0）

计算模型整体参数量

num_parameters = sum(torch.numel(parameter) for parameter in model.parameters())

导入另一个模型的相同部分到新的模型

模型导入参数时，如果两个模型结构不一致，则直接导入模型参数会报错，用下面方法可以把另一个模型的相同部分导入新的模型中

def load_model(model, model_path, optimizer=None, resume=False, 
               lr=None, lr_step=None):
  start_epoch = 0
  checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
  print('loaded {}, epoch {}'.format(model_path, checkpoint['epoch']))
  state_dict_ = checkpoint['state_dict']
  state_dict = {}
  
  # convert data_parallal to model
  for k in state_dict_:
    if k.startswith('module') and not k.startswith('module_list'):
      state_dict[k[7:]] = state_dict_[k]
    else:
      state_dict[k] = state_dict_[k]
  model_state_dict = model.state_dict()

  # check loaded parameters and created model parameters
  msg = 'If you see this, your model does not fully load the ' + \
        'pre-trained weight. Please make sure ' + \
        'you have correctly specified --arch xxx ' + \
        'or set the correct --num_classes for your own dataset.'
  for k in state_dict:
    if k in model_state_dict:
      if state_dict[k].shape != model_state_dict[k].shape:
        print('Skip loading parameter {}, required shape{}, '\
              'loaded shape{}. {}'.format(
          k, model_state_dict[k].shape, state_dict[k].shape, msg))
        state_dict[k] = model_state_dict[k]
    else:
      print('Drop parameter {}.'.format(k) + msg)
  for k in model_state_dict:
    if not (k in state_dict):
      print('No param {}.'.format(k) + msg)
      state_dict[k] = model_state_dict[k]
  model.load_state_dict(state_dict, strict=False)

  # resume optimizer parameters
  if optimizer is not None and resume:
    if 'optimizer' in checkpoint:
      optimizer.load_state_dict(checkpoint['optimizer'])
      start_epoch = checkpoint['epoch']
      start_lr = lr
      for step in lr_step:
        if start_epoch >= step:
          start_lr *= 0.1
      for param_group in optimizer.param_groups:
        param_group['lr'] = start_lr
      print('Resumed optimizer with start lr', start_lr)
    else:
      print('No optimizer parameters in checkpoint.')
  if optimizer is not None:
    return model, optimizer, start_epoch
  else:
    return model

其他注意事项

1. 不要使用太大的线性层，因为nn.Linear(M,N)使用的是O(mn)的内存，线性层太大很容易超出现有的显存
2. 不要在太长的序列上使用RNN，因为RNN反向传播使用的BPTT算法，其需要的内存和输入序列的长度呈线性关系
3. model(x)前用model.train()和model.eval()切换网络状态
4. 不要计算梯度的代码用with torch.no_grad()包含起来
5. model.eval()和torch.no_grad()的区别在于，**model.eval()**是将网络切换为测试状态，例如BN和dropout在训练和测试阶段使用不同的计算方法。torch.no_grad()是关闭PyTorch张量的自动求导机制，以减少存储使用和加速计算，得到的结果无法进行loss.backward()
6. model.zero_grad()会把整个模型的参数的梯度都归零，而optimizer.zero_grad()只会把传入其中的参数的梯度归零
7. loss.backward()前用optimizer.zero_grad()清除累计梯度
8. torch.utils.data.DataLoader中尽量设置pin_memory=True, 对特别小的数据集MNIST设置pin_memory=False反而更快有些。num_workers的设置需要在实验中找到最快的取值
9. 用del及时删除不用的中间变量，节约GPU存储
10. 使用inplace操作可节约GPU存储
11. 减少CPU和GPU之间的数据传输
12. 使用半精度浮点数half()会有一定的速度提升，具体效率依赖于GPU型号，需要小心数值精度过低带来的稳定性问题
13. 时常使用assert tensor.size() == (N,D,H,W)作为调试手段，确保张量维度和设想中一致
14. 除了标记y外，尽量少使用一维张量，使用n*1的二维张量替代
15. 统计代码各部分耗时

with torch.autograd.profiler.profile(enabled=True, use_cuda=False) as profile:
    ...
print(profile)

# 或者在命令行运行
python -m torch.utils.bottleneck main.py

16. 使用TorchSnooper来调试PyTorch代码，程序在执行时候，就会自动print出来每一行的执行结果的tensor形状，数据类型，设备，是否需要梯度的信息

# pip install torchsnooper
import torchsnooper

# 对于函数，使用修饰器
@torchsnooper.snoop()

# 如果不是函数，使用 with 语句来激活 TorchSnooper，把训练的那个循环装进 with 语句中去。
with torchsnooper.snoop():
    原本的代码