【转载】PyTorch 实用技巧

https://github.com/zxdefying/pytorch_tricks

在github上看到一篇关于PyTorch实战技巧的文章，感觉很棒，感恩作者，为防止迷路，转载一下先。

1、指定GPU编号

设置当前使用的GPU设备仅为0号设备，设备名称为 /gpu:0：os.environ[“CUDA_VISIBLE_DEVICES”] = “0”
设置当前使用的GPU设备为0,1号两个设备，名称依次为 /gpu:0、/gpu:1： os.environ[“CUDA_VISIBLE_DEVICES”] = “0,1” ，根据顺序表示优先使用0号设备,然后使用1号设备。
指定GPU的命令需要放在和神经网络相关的一系列操作的前面。

2、查看模型每层输出详情

Keras有一个简洁的API来查看模型的每一层输出尺寸，这在调试网络时非常有用。现在在PyTorch中也可以实现这个功能。

使用很简单，如下用法：

from torchsummary import summary
summary(your_model, input_size=(channels, H, W))

input_size 是根据你自己的网络模型的输入尺寸进行设置。

3、梯度裁剪（Gradient Clipping）

import torch.nn as nn

outputs = model(data)
loss= loss_fn(outputs, target)
optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), max_norm=20, norm_type=2)
optimizer.step()

nn.utils.clip_grad_norm_ 的参数：

parameters – 一个基于变量的迭代器，会进行梯度归一化
max_norm – 梯度的最大范数
norm_type – 规定范数的类型，默认为L2
不椭的椭圆 提出：梯度裁剪在某些任务上会额外消耗大量的计算时间，可移步评论区查看详情。

4、扩展单张图片维度

因为在训练时的数据维度一般都是 (batch_size, c, h, w)，而在测试时只输入一张图片，所以需要扩展维度，扩展维度有多个方法：

import cv2
import torch

image = cv2.imread(img_path)
image = torch.tensor(image)
print(image.size())

img = image.view(1, *image.size())
print(img.size())

# output:
# torch.Size([h, w, c])
# torch.Size([1, h, w, c])

或

import cv2
import numpy as np

image = cv2.imread(img_path)
print(image.shape)
img = image[np.newaxis, :, :, :]
print(img.shape)

# output:
# (h, w, c)
# (1, h, w, c)

import cv2
import torch

image = cv2.imread(img_path)
image = torch.tensor(image)
print(image.size())

img = image.unsqueeze(dim=0)  
print(img.size())

img = img.squeeze(dim=0)
print(img.size())

# output:
# torch.Size([(h, w, c)])
# torch.Size([1, h, w, c])
# torch.Size([h, w, c])

tensor.unsqueeze(dim)：扩展维度，dim指定扩展哪个维度。

tensor.squeeze(dim)：去除dim指定的且size为1的维度，维度大于1时，squeeze()不起作用，不指定dim时，去除所有size为1的维度。

5、独热编码

在PyTorch中使用交叉熵损失函数的时候会自动把label转化成onehot，所以不用手动转化，而使用MSE需要手动转化成onehot编码。

import torch
class_num = 8
batch_size = 4

def one_hot(label):
    """
    将一维列表转换为独热编码
    """
    label = label.resize_(batch_size, 1)
    m_zeros = torch.zeros(batch_size, class_num)
    # 从 value 中取值，然后根据 dim 和 index 给相应位置赋值
    onehot = m_zeros.scatter_(1, label, 1)  # (dim,index,value)

    return onehot.numpy()  # Tensor -> Numpy

label = torch.LongTensor(batch_size).random_() % class_num  # 对随机数取余
print(one_hot(label))

# output:
[[0. 0. 0. 1. 0. 0. 0. 0.]
 [0. 0. 0. 0. 1. 0. 0. 0.]
 [0. 0. 1. 0. 0. 0. 0. 0.]
 [0. 1. 0. 0. 0. 0. 0. 0.]]

​

6、防止验证模型时爆显存

验证模型时不需要求导，即不需要梯度计算，关闭autograd，可以提高速度，节约内存。如果不关闭可能会爆显存。

with torch.no_grad():
    # 使用model进行预测的代码
    pass

Pytorch 训练时无用的临时变量可能会越来越多，导致 out of memory ，可以使用下面语句来清理这些不需要的变量。

官网 上的解释为：

Releases all unoccupied cached memory currently held by the caching allocator so that those can be used in other GPU application and visible innvidia-smi. torch.cuda.empty_cache()

意思就是PyTorch的缓存分配器会事先分配一些固定的显存，即使实际上tensors并没有使用完这些显存，这些显存也不能被其他应用使用。这个分配过程由第一次CUDA内存访问触发的。

而 torch.cuda.empty_cache() 的作用就是释放缓存分配器当前持有的且未占用的缓存显存，以便这些显存可以被其他GPU应用程序中使用，并且通过 nvidia-smi命令可见。注意使用此命令不会释放tensors占用的显存。

对于不用的数据变量，Pytorch 可以自动进行回收从而释放相应的显存。

更详细的优化可以查看 优化显存使用 和 显存利用问题。

7、学习率衰减

import torch.optim as optim
from torch.optim import lr_scheduler

# 训练前的初始化
optimizer = optim.Adam(net.parameters(), lr=0.001)
scheduler = lr_scheduler.StepLR(optimizer, 10, 0.1)  # # 每过10个epoch，学习率乘以0.1

# 训练过程中
for n in n_epoch:
    scheduler.step()
    ...

8、冻结某些层的参数

参考：Pytorch 冻结预训练模型的某一层

在加载预训练模型的时候，我们有时想冻结前面几层，使其参数在训练过程中不发生变化。

我们需要先知道每一层的名字，通过如下代码打印：

net = Network()  # 获取自定义网络结构
for name, value in net.named_parameters():
    print('name: {0},\t grad: {1}'.format(name, value.requires_grad))

假设前几层信息如下：

name: cnn.VGG_16.convolution1_1.weight,	 grad: True
name: cnn.VGG_16.convolution1_1.bias,	 grad: True
name: cnn.VGG_16.convolution1_2.weight,	 grad: True
name: cnn.VGG_16.convolution1_2.bias,	 grad: True
name: cnn.VGG_16.convolution2_1.weight,	 grad: True
name: cnn.VGG_16.convolution2_1.bias,	 grad: True
name: cnn.VGG_16.convolution2_2.weight,	 grad: True
name: cnn.VGG_16.convolution2_2.bias,	 grad: True

后面的True表示该层的参数可训练，然后我们定义一个要冻结的层的列表：

no_grad = [
    'cnn.VGG_16.convolution1_1.weight',
    'cnn.VGG_16.convolution1_1.bias',
    'cnn.VGG_16.convolution1_2.weight',
    'cnn.VGG_16.convolution1_2.bias'
]

冻结方法如下：

net = Net.CTPN()  # 获取网络结构
for name, value in net.named_parameters():
    if name in no_grad:
        value.requires_grad = False
    else:
        value.requires_grad = True

冻结后我们再打印每层的信息：

name: cnn.VGG_16.convolution1_1.weight,	 grad: False
name: cnn.VGG_16.convolution1_1.bias,	 grad: False
name: cnn.VGG_16.convolution1_2.weight,	 grad: False
name: cnn.VGG_16.convolution1_2.bias,	 grad: False
name: cnn.VGG_16.convolution2_1.weight,	 grad: True
name: cnn.VGG_16.convolution2_1.bias,	 grad: True
name: cnn.VGG_16.convolution2_2.weight,	 grad: True
name: cnn.VGG_16.convolution2_2.bias,	 grad: True

可以看到前两层的weight和bias的requires_grad都为False，表示它们不可训练。

最后在定义优化器时，只对requires_grad为True的层的参数进行更新。

optimizer = optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=0.01)

9、对不同层使用不同学习率
我们对模型的不同层使用不同的学习率。

还是使用这个模型作为例子：

net = Network()  # 获取自定义网络结构
for name, value in net.named_parameters():
    print('name: {}'.format(name))

# 输出：
# name: cnn.VGG_16.convolution1_1.weight
# name: cnn.VGG_16.convolution1_1.bias
# name: cnn.VGG_16.convolution1_2.weight
# name: cnn.VGG_16.convolution1_2.bias
# name: cnn.VGG_16.convolution2_1.weight
# name: cnn.VGG_16.convolution2_1.bias
# name: cnn.VGG_16.convolution2_2.weight
# name: cnn.VGG_16.convolution2_2.bias

对 convolution1 和 convolution2 设置不同的学习率，首先将它们分开，即放到不同的列表里：

conv1_params = []
conv2_params = []

for name, parms in net.named_parameters():
    if "convolution1" in name:
        conv1_params += [parms]
    else:
        conv2_params += [parms]

# 然后在优化器中进行如下操作：
optimizer = optim.Adam(
    [
        {"params": conv1_params, 'lr': 0.01},
        {"params": conv2_params, 'lr': 0.001},
    ],
    weight_decay=1e-3,
)

我们将模型划分为两部分，存放到一个列表里，每部分就对应上面的一个字典，在字典里设置不同的学习率。当这两部分有相同的其他参数时，就将该参数放到列表外面作为全局参数，如上面的weight_decay。

也可以在列表外设置一个全局学习率，当各部分字典里设置了局部学习率时，就使用该学习率，否则就使用列表外的全局学习率。

10、模型相关操作

10.1、保存加载模型基本用法

1、保存加载整个模型

保存整个网络模型（网络结构+权重参数）。

torch.save(model, 'net.pkl')

直接加载整个网络模型（可能比较耗时）。

model = torch.load('net.pkl')

2、只保存加载模型参数

只保存模型的权重参数（速度快，占内存少）。

torch.save(model.state_dict(), 'net_params.pkl')

因为我们只保存了模型的参数，所以需要先定义一个网络对象，然后再加载模型参数。

# 构建一个网络结构
model = ClassNet()
# 将模型参数加载到新模型中
state_dict = torch.load('net_params.pkl')
model.load_state_dict(state_dict)

保存模型进行推理测试时，只需保存训练好的模型的权重参数，即推荐第二种方法。

主要用法就是上面这些，接下来讲一下PyTorch中保存加载模型内部的一些原理，以及我们可能会遇到的一些特殊的需求。

10.2、保存加载自定义模型

上面保存加载的 net.pkl 其实一个字典，通常包含如下内容：

网络结构：输入尺寸、输出尺寸以及隐藏层信息，以便能够在加载时重建模型。
模型的权重参数：包含各网络层训练后的可学习参数，可以在模型实例上调用 state_dict() 方法来获取，比如前面介绍只保存模型权重参数时用到的 model.state_dict()。
优化器参数：有时保存模型的参数需要稍后接着训练，那么就必须保存优化器的状态和所其使用的超参数，也是在优化器实例上调用 state_dict() 方法来获取这些参数。
其他信息：有时我们需要保存一些其他的信息，比如 epoch，batch_size 等超参数。

知道了这些，那么我们就可以自定义需要保存的内容，比如：

# saving a checkpoint assuming the network class named ClassNet
checkpoint = {'model': ClassNet(),
              'model_state_dict': model.state_dict(),
              'optimizer_state_dict': optimizer.state_dict(),
              'epoch': epoch}
​
torch.save(checkpoint, 'checkpoint.pkl')

上面的 checkpoint 是个字典，里面有4个键值对，分别表示网络模型的不同信息。

然后我们要加载上面保存的自定义的模型：

def load_checkpoint(filepath):
    checkpoint = torch.load(filepath)
    model = checkpoint['model']  # 提取网络结构
    model.load_state_dict(checkpoint['model_state_dict'])  # 加载网络权重参数
    optimizer = TheOptimizerClass()
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])  # 加载优化器参数
    
    for parameter in model.parameters():
        parameter.requires_grad = False
    model.eval()
    
    return model
    
model = load_checkpoint('checkpoint.pkl')

如果加载模型只是为了进行推理测试，则将每一层的 requires_grad 置为 False，即固定这些权重参数；还需要调用 model.eval() 将模型置为测试模式，主要是将 dropout 和 batch normalization 层进行固定，否则模型的预测结果每次都会不同。

如果希望继续训练，则调用 model.train()，以确保网络模型处于训练模式。

state_dict() 也是一个Python字典对象，model.state_dict() 将每一层的可学习参数映射为参数矩阵，其中只包含具有可学习参数的层(卷积层、全连接层等)。

比如下面这个例子：

# Define model
class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.conv1 = nn.Conv2d(3, 8, 5)
        self.bn = nn.BatchNorm2d(8)
        self.conv2 = nn.Conv2d(8, 16, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 10)
​
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.bn(x)
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
    
    # Initialize model
    model = TheModelClass()
​
    # Initialize optimizer
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
​
    print("Model's state_dict:")
    for param_tensor in model.state_dict():
        print(param_tensor, "\t", model.state_dict()[param_tensor].size())
​
    print("Optimizer's state_dict:")
    for var_name in optimizer.state_dict():
        print(var_name, "\t", optimizer.state_dict()[var_name])
输出为：

Model's state_dict:
conv1.weight            torch.Size([8, 3, 5, 5])
conv1.bias              torch.Size([8])
bn.weight               torch.Size([8])
bn.bias                 torch.Size([8])
bn.running_mean         torch.Size([8])
bn.running_var          torch.Size([8])
bn.num_batches_tracked  torch.Size([])
conv2.weight            torch.Size([16, 8, 5, 5])
conv2.bias              torch.Size([16])
fc1.weight              torch.Size([120, 400])
fc1.bias                torch.Size([120])
fc2.weight              torch.Size([10, 120])
fc2.bias                torch.Size([10])
Optimizer's state_dict:
state            {}
param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [139805696932024, 139805483616008, 139805483616080, 139805483616152, 139805483616440, 139805483616512, 139805483616584, 139805483616656, 139805483616728, 139805483616800]}]

可以看到 model.state_dict() 保存了卷积层，BatchNorm层和最大池化层的信息；而 optimizer.state_dict() 则保存的优化器的状态和相关的超参数。

10.3、跨设备保存加载模型

1、在 CPU 上加载在 GPU 上训练并保存的模型（Save on GPU, Load on CPU）：

device = torch.device('cpu')
model = TheModelClass()
# Load all tensors onto the CPU device
model.load_state_dict(torch.load('net_params.pkl', map_location=device))
map_location：a function, torch.device, string or a dict specifying how to remap storage locations

令 torch.load() 函数的 map_location 参数等于 torch.device(‘cpu’) 即可。 这里令 map_location 参数等于 ‘cpu’ 也同样可以。

2、在 GPU 上加载在 GPU 上训练并保存的模型（Save on GPU, Load on GPU）：

device = torch.device("cuda")
model = TheModelClass()
model.load_state_dict(torch.load('net_params.pkl'))
model.to(device)

在这里使用 map_location 参数不起作用，要使用 model.to(torch.device(“cuda”)) 将模型转换为CUDA优化的模型。

还需要对将要输入模型的数据调用 data = data.to(device)，即将数据从CPU转移到GPU。请注意，调用 my_tensor.to(device) 会返回一个 my_tensor 在 GPU 上的副本，它不会覆盖 my_tensor。因此需要手动覆盖张量：my_tensor = my_tensor.to(device)。

3、在 GPU 上加载在 GPU 上训练并保存的模型（Save on CPU, Load on GPU）

device = torch.device("cuda")
model = TheModelClass()
model.load_state_dict(torch.load('net_params.pkl', map_location="cuda:0"))
model.to(device)

当加载包含GPU tensors的模型时，这些tensors 会被默认加载到GPU上，不过是同一个GPU设备。

当有多个GPU设备时，可以通过将 map_location 设定为 cuda:device_id 来指定使用哪一个GPU设备，上面例子是指定编号为0的GPU设备。

其实也可以将 torch.device(“cuda”) 改为 torch.device(“cuda:0”) 来指定编号为0的GPU设备。

最后调用 model.to(torch.device(‘cuda’)) 来将模型的tensors转换为 CUDA tensors。

下面是PyTorch官方文档上的用法，可以进行参考：

>>> torch.load('tensors.pt')
# Load all tensors onto the CPU
>>> torch.load('tensors.pt', map_location=torch.device('cpu'))
# Load all tensors onto the CPU, using a function
>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage)
# Load all tensors onto GPU 1
>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
# Map tensors from GPU 1 to GPU 0
>>> torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})

10.4、CUDA 的用法

在PyTorch中和GPU相关的几个函数：

import torch

# 判断cuda是否可用；
print(torch.cuda.is_available())

# 获取gpu数量；
print(torch.cuda.device_count())

# 获取gpu名字；
print(torch.cuda.get_device_name(0))

# 返回当前gpu设备索引，默认从0开始；
print(torch.cuda.current_device())

我的电脑输出为：

True
1
GeForce RTX 2080 Ti
0

有时我们需要把数据和模型从cpu移到gpu中，有以下两种方法：

use_cuda = torch.cuda.is_available()

# 方法一：
if use_cuda:
    data = data.cuda()
    model.cuda()

# 方法二：
device = torch.device("cuda" if use_cuda else "cpu")
data = data.to(device)
model.to(device)

个人比较习惯第二种方法，可以少一个 if 语句。而且该方法还可以通过设备号指定使用哪个GPU设备，比如使用0号设备：

device = torch.device("cuda:0" if use_cuda else "cpu")

11、Pytorch内置one_hot函数

感谢 yangyangyang 补充：Pytorch 1.1后，one_hot可以直接用torch.nn.functional.one_hot。
然后我将Pytorch升级到1.2版本，试用了下 one_hot 函数，确实很方便。
具体用法如下：

import torch.nn.functional as F
import torch

tensor =  torch.arange(0, 5) % 3  # tensor([0, 1, 2, 0, 1])
one_hot = F.one_hot(tensor)

# 输出：
# tensor([[1, 0, 0],
#         [0, 1, 0],
#         [0, 0, 1],
#         [1, 0, 0],
#         [0, 1, 0]])

F.one_hot会自己检测不同类别个数，生成对应独热编码。我们也可以自己指定类别数：

tensor =  torch.arange(0, 5) % 3  # tensor([0, 1, 2, 0, 1])
one_hot = F.one_hot(tensor, num_classes=5)

# 输出：
# tensor([[1, 0, 0, 0, 0],
#         [0, 1, 0, 0, 0],
#         [0, 0, 1, 0, 0],
#         [1, 0, 0, 0, 0],
#         [0, 1, 0, 0, 0]])

升级 Pytorch (cpu版本)的命令：conda install pytorch torchvision -c pytorch