PyTorch：模型训练、评估和预测

-柚子皮-

模型训练、评估和预测

模型训练

pytorch可以给我们提供两种方式来切换训练和评估(推断)的模式，分别是：model.train() 和 model.eval()。一般用法是：在训练开始之前写上 model.trian() ，在测试时写上 model.eval() 。

model.train()作用是 启用 batch normalization 和 dropout 。保证 BN 层能够用到 每一批数据 的均值和方差。对于 Dropout，model.train() 是 随机取一部分 网络连接来训练更新参数。

单机训练

传统的batch训练函数

简单的说就是进来一个batch的数据，计算一次梯度，更新一次网络。

optimizer = Adam(model.parameters(), lr=args.lr)

for i,(images,target) in enumerate(train_loader):
    # 1. input output
    images = images.cuda(non_blocking=True)
    target = torch.from_numpy(np.array(target)).float().cuda(non_blocking=True)
    outputs = model(images)
    loss = criterion(outputs,target)

    # 2. backward
    optimizer.zero_grad()   # reset gradient
    loss.backward()
    optimizer.step()            

1 获取loss：输入图像和标签，通过infer计算得到预测值，计算损失函数；
2 optimizer.zero_grad() 清空过往梯度；  lz:这里需要手动进行optimizer.zero_grad()，不然pytorch中默认是累加的。

另外，根据optimizer的定义，也可以使用model.zero_grad() 。因为当仅有一个model，同时optimizer只包含这一个model的参数，那么model.zero_grad和optimizer.zero_grad没有区别，可以任意使用。[optimizer.zero_grad和model.zero_grad有啥区别？ - 知乎]

optimizer = optim.Adam(list(model1.parameters())+list(model2.parameters()), lr=0.01)
3 loss.backward() 反向传播，计算当前梯度；

4 当然在更新参数之前可以加一个梯度裁剪：nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm=1.0)
5 optimizer.step() 根据梯度更新网络参数

示例

def compile(self, optimizer_name, loss_func_name=None, metrics_name=None):
    """
    :param optimizer_name: String (name of optimizer) or optimizer instance. See [optimizers](https://pytorch.org/docs/stable/optim.html).
    :param loss_func_name: String (name of objective function) or objective function. See [losses](https://pytorch.org/docs/stable/nn.functional.html#loss-functions).
    :param metrics_name: List of metrics to be evaluated by the model during training and testing. Typically you will use `metrics=['accuracy']`.
    """
    self.optimizer = self._get_optim(optimizer_name)
    self.loss_func = self._get_loss_func(loss_func_name)
    self.metrics = self._get_metrics(metrics_name)

def _get_optim(self, optimizer):
    if isinstance(optimizer, str):
        if optimizer == "sgd":
            optim = torch.optim.SGD(self.parameters(), lr=0.01)
        elif optimizer == "adam":
            optim = torch.optim.Adam(self.parameters())
        elif optimizer == "adagrad":
            optim = torch.optim.Adagrad(self.parameters())
        elif optimizer == "rmsprop":
            optim = torch.optim.RMSprop(self.parameters())
        else:
            raise NotImplementedError
    else:
        optim = optimizer
    return optim

def optimize(self, max_norm=None, norm_type=2):
    '''
    进行最优化操作
    '''
    with time_block("backward_in_optimize"):
        self.loss.backward()
    if max_norm:
        torch.nn.utils.clip_grad_norm_(self.parameters(), max_norm, norm_type)
    print("next(self.parameters()).device: ", next(self.parameters()).device)
    self.optimizer.step()

init中初始化时：

self.to(device)

self.compile(optimizer, loss, metrics) #需要在self.to(device)后面，因为optim = torch.optim.Adagrad(self.parameters())中next(self.parameters()).device的类型还未确定，self.parameters()为gpu时，optim才是gpu中，self.optimizer.step()执行时才不会出错：RuntimeError: Expected object of backend CPU but got backend CUDA for argument #4 'tensor1'。

使用梯度累加的batch训练函数

梯度累加就是，每次获取1个batch的数据，计算1次梯度，梯度不清空，不断累加，累加一定次数后，根据累加的梯度更新网络参数，然后清空梯度，进行下一次循环。

for i,(images,target) in enumerate(train_loader):
    # 1. input output
    images = images.cuda(non_blocking=True)
    target = torch.from_numpy(np.array(target)).float().cuda(non_blocking=True)
    outputs = model(images)
    loss = criterion(outputs,target)

    # 2.1 loss regularization
    loss = loss/accumulation_steps   
    # 2.2 back propagation
    loss.backward()
    # 3. update parameters of net
    if((i+1)%accumulation_steps)==0:
        # optimizer the net
        optimizer.step()        # update parameters of net
        optimizer.zero_grad()   # reset gradient

1 获取loss：输入图像和标签，通过infer计算得到预测值，计算损失函数；
2 loss.backward() 反向传播，计算当前梯度；   lz:2和3之间应该有一个pytorch自动累加梯度的过程（应该就是backward里面的过程吧）。
3 多次循环步骤1-2，不清空梯度，使梯度累加在已有梯度上；
4 梯度累加了一定次数后，先optimizer.step() 根据累计的梯度更新网络参数，然后optimizer.zero_grad() 清空过往梯度，为下一波梯度累加做准备；

与非梯度累加区别

      一定条件下，batchsize越大训练效果越好，梯度累加则实现了batchsize的变相扩大，如果accumulation_steps为8，则batchsize '变相' 扩大了8倍，是解决显存受限的一个不错的trick。

使用时需要注意，学习率也要适当放大。经验法则是，如果batch size加倍，那么学习率就加倍。

PyTorch中在反向传播前为什么要手动将梯度清零？

有时候会不会觉得optimizer.zero_grad()这句很多余，怎么不自动清零，需要手动清零呢？原因在于在PyTorch中，计算得到的梯度值会自动进行累加（而不是覆盖）。pytorch选择自动累加而不是覆盖的原因：

1 从内存消耗的角度来看。这种模式可以让梯度玩出更多花样，比如说上面讲到的梯度累加（gradient accumulation）实现的“显存受限”解决：在内存大小不够的情况下叠加多个batch的grad作为一个大batch进行迭代，因为二者得到的梯度是等价的综上可知，这种梯度累加的思路是对内存的极大友好，是由FAIR的设计理念出发的。即当你GPU显存较少时，你又想要调大batch-size，此时你就可以利用PyTorch的这个性质进行梯度的累加来进行backward。

示例：

在你已经达到计算资源上限的情况下，你的batch size仍然太小(比如8)，然后我们需要模拟一个更大的batch size来进行梯度下降，以提供一个良好的估计。

假设我们想要达到128的batch size大小。我们需要以batch size为8执行16个前向传播和向后传播，然后再执行一次优化步骤。

# clear last step
optimizer.zero_grad()

# 16 accumulated gradient steps
scaled_loss = 0
for accumulated_step_i in range(16):
     out = model.forward()
     loss = some_loss(out,y)    
     loss.backward()
      scaled_loss += loss.item()
      
# update weights after 8 steps. effective batch = 8*16
optimizer.step()

# loss is now scaled up by the number of accumulated batches
actual_loss = scaled_loss / 16

2 利用梯度累加，可以在最多保存一张计算图的情况下进行multi-task任务的训练：

从PyTorch的设计原理上来说，在每次进行前向计算得到pred时，会产生一个用于梯度回传的计算图，这张图储存了进行back propagation需要的中间结果，当调用了.backward()后，会从内存中将这张图进行释放。

for idx, data in enumerate(train_loader):
    xs, ys = data
    optmizer.zero_grad()
    # 计算d(l1)/d(x)
    pred1 = model1(xs) #生成graph1
    loss1 = loss_fn1(pred1, ys)
    loss1.backward()  #释放graph1

    # 计算d(l2)/d(x)
    pred2 = model2(xs)#生成graph2
    loss2 = loss_fn2(pred2, ys)
    loss2.backward()  #释放graph2

    # 使用d(l1)/d(x)+d(l2)/d(x)进行优化
    optmizer.step()

[PyTorch中在反向传播前为什么要手动将梯度清零？]

模型训练加速：16-bit 精度

16bit精度是将内存占用减半的惊人技术。大多数模型使用32bit精度数字进行训练。然而，最近的研究发现，16bit模型也可以工作得很好。混合精度意味着对某些内容使用16bit，但将权重等内容保持在32bit。

要在Pytorch中使用16bit精度，请安装NVIDIA的apex库，并对你的模型进行这些更改。

# enable 16-bit on the model and the optimizer
model, optimizers = amp.initialize(model, optimizers, opt_level='O2')

# when doing .backward, let amp do it so it can scale the loss
with amp.scale_loss(loss, optimizer) as scaled_loss:                      
    scaled_loss.backward()

amp包会处理好大部分事情。如果梯度爆炸或趋向于0，它甚至会缩放loss。

在lightning中，启用16bit并不需要修改模型中的任何内容，也不需要执行我上面所写的操作。设置Trainer(precision=16)就可以了。

trainer = Trainer(amp_level='O2', use_amp=False)
trainer.fit(model)

Gpu训练

使用 PyTorch 查看 CUDA 和 cuDNN 版本

import torch
print(torch.__version__)

print(torch.version.cuda)
print(torch.backends.cudnn.version())

2.0.1+cu117
11.7
8600

注意低版本的 pytorch 是否支持更高版本的 cuda。（高版本的pytorch一般能兼容低版本cuda）

pytorch和cudatoolkit版本对应关系：[Previous PyTorch Versions | PyTorch]

2.cudatoolkit版本和系统cuda对应关系：

https://www.zhihu.com/question/344950161/answer/818139888

3.系统cuda和nvidia对应关系：

https://blog.csdn.net/He_9520/article/details/100032803

4.cuda和cuDNN的关系和对应关系：

https://www.jianshu.com/p/622f47f94784

https://www.cnblogs.com/yeran/p/11345990.html

bugs

1 cuda版本和pytorch版本对应不上可能会出错：如RuntimeError: cuDNN version incompatibility: PyTorch was compiled  against (8, 5, 0) but found runtime version (8, 2, 0)

2 要进行2个张量的运算，就必须都在CPU或者都在GPU上。对于不同存储位置的变量，我们是不可以对他们直接进行计算的。存储在不同位置中的数据是不可以直接进行交互计算的。否则出错：RuntimeError: Expected object of backend CUDA but got backend CPU for ***。[PyTorch 20.GPU训练]

判断模型和数据在哪里

判断模型是在CPU还是GPU上
model = nn.LSTM(input_size=10, hidden_size=4, num_layers=1, batch_first=True)
print(next(model.parameters()).device)  # 输出：cpu
model = model.cuda()
print(next(model.parameters()).device)  # 输出：cuda:0
model = model.cpu()
print(next(model.parameters()).device)  # 输出：cpu

判断数据是在CPU还是GPU上
data = torch.ones([2, 3])
print(data.device)  # 输出：cpu 
data = data.cuda()
print(data.device)  # 输出：cuda:0 
data = data.cpu()
print(data.device)  # 输出：cpu

把数据和模型从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")

[PyTorch：tensor-基本操作---CUDA 的用法]

1 如果模型已经在GPU上了，model.to(device)/model.cuda()不会做任何事情。

2 对于自己创建的模型类，由于继承了 torch.nn.Module ，则可同样使用 model.to(device)/model.cuda() 来将模型中用到的所有参数都存储到显存中去。

model_cuda = model.cuda()后，将 cpu类型数据 投入 model_cuda 中去可以发现，系统会报错，而将 .cuda 之后的gpu数据投入 model_cuda 才能正常工作，并且输出的也是具有cuda的数据类型。

self.dnn_layer = MLP(input_dim, dnn_hidden_units, activation=dnn_activation,
                    use_bn=dnn_use_bn, l2_reg=l2_reg_dnn, dropout_rate=dnn_dropout,
                    init_std=init_std, device=device)
self.to(device)
...
print("self.dnn_layer.device: ", self.dnn_layer.device) #cpu，模型to('gpu')时函数本身没有变化，只是参数put到gpu了。
print("self.dnn_layer.parameters()).device: ", next(self.dnn_layer.parameters()).device) #gpu

3 当我们对模型存储到显存中去之后，那么这个模型中的方法后面所创建出来的Tensor是不是都会默认变成cuda的数据类型？答案是否定的。

[浅谈将Pytorch模型从CPU转换成GPU]

指定GPU编号

os.environ['CUDA_VISIBLE_DEVICES']是设定程序对哪几张卡可视。指定GPU的命令需要放在和神经网络相关的一系列操作的前面。

1 设置当前使用的GPU设备仅为0号设备，设备名称为 /gpu:0：

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

2 设置当前使用的GPU设备为0,1号两个设备，名称依次为 /gpu:0、/gpu:1，根据顺序表示优先使用0号设备,然后使用1号设备：

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

3 一般设定成功之后，接下来程序中任何有关卡号的指定都是相对的：

os.environ['CUDA_VISIBLE_DEVICES'] = '1,2,3'，生效之后，再设置torch.cuda.set_device(0)，此时pytorch将会使用1号cuda.

出错：RuntimeError: CUDA error: out of memory?

1 先用nvidia查看gpu内存占用情况，可能是gpu0正被别人在用且无内存了

这时可以通过os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"指定其它可用gpu
2 batch_size过大

调小点再试试

单机多GPU训练

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"只能指定多个gpu可用，但是只能用其中一个。要用多个可以这样：

args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")

if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs")
    # dim = 0 [64, xxx] -> [32, ...], [32, ...] on 2GPUs
    model = torch.nn.DataParallel(model)
model = model.to(args.device)

但是有些代码还是可能出错，比如

[Caught StopIteration in replica 1 on device 1 error while Training on GPU]

多节点GPU训练，在单个节点上多GPU更快的训练distributedDataParallel

[移动到多个GPUs中]

单机多GPU的实现代码

[Pytorch 高效使用GPU的操作]

分布式训练

[PyTorch：模型训练-分布式训练]

训练设置

梯度裁剪

nn.utils.clip_grad_norm_ 的参数
parameters – 一个基于变量的迭代器，会进行梯度归一化
max_norm – 梯度的最大范数
norm_type – 规定范数的类型，默认为L2
Note: 梯度裁剪在某些任务上会额外消耗大量的计算时间。

[torch.nn.utils.clip_grad_norm_ — PyTorch master documentation]

示例

def optimize(self, max_norm=None, norm_type=2):
        '''
        进行最优化操作
        '''
        self.loss.backward()
        if max_norm:
            torch.nn.utils.clip_grad_norm_(self.parameters(), max_norm, norm_type)
        self.optimizer.step()

import torch.nn as nn

outputs = model(data)
loss= loss_fn(outputs, target)
optimizer.zero_grad()

self.optimize()

16-bit 精度

16bit精度是将内存占用减半的惊人技术。大多数模型使用32bit精度数字进行训练。然而，最近的研究发现，16bit模型也可以工作得很好。混合精度意味着对某些内容使用16bit，但将权重等内容保持在32bit。

要在Pytorch中使用16bit精度，请安装NVIDIA的apex库，并对你的模型进行这些更改。

# enable 16-bit on the model and the optimizer
model, optimizers = amp.initialize(model, optimizers, opt_level='O2')

# when doing .backward, let amp do it so it can scale the loss
with amp.scale_loss(loss, optimizer) as scaled_loss:                      
    scaled_loss.backward()

amp包会处理好大部分事情。如果梯度爆炸或趋向于0，它甚至会缩放loss。

防止验证模型时爆显存torch.no_grad()

torch.no_grad()是一个上下文管理器，用来禁止梯度的计算，通常用来网络推断中（因为验证模型时不需要求导，即不需要梯度计算），关闭autograd可以减少计算内存的使用量，提高速度，如果不关闭可能会爆显存。
with torch.no_grad():
    # 使用model进行预测的代码
    pass

示例
>>> a = torch.tensor([1.0, 2.0], requires_grad=True)
>>> with torch.no_grad():
...     b = n.pow(2).sum()
...
>>> b
tensor(5.)
>>> b.requires_grad
False
>>> c = a.pow(2).sum()
>>> c.requires_grad
True
上面的例子中，

当a的requires_grad=True时，不使用torch.no_grad()，c.requires_grad为True；

使用torch.no_grad()时，b.requires_grad为False，当不需要进行反向传播时（推断）或不需要计算梯度（网络输入）时，requires_grad=True会占用更多的计算资源及存储资源。

Note: 需要注意的是，torch.no_grad()只是禁止梯度计算，将计算过程中的tensor的requires_grad设置为False，但不会将模型参数如weight这样的tensor的requires_grad设置为False，因为如果训练和测试同时进行，跳出predict的torch.no_grad()后的训练过程还是需要weight中requires_grad=True的。

释放保留的计算图

一个最简单撑爆你的内存的方法是为了记录日志存储你的loss。

losses = []
...
losses.append(loss)

print(f'current loss: {torch.mean(losses)'})
上面的问题是，loss仍然包含有整个图的副本。在这种情况下，调用.item()来释放它。

# good
losses.append(loss.item())

torch.cuda.empty_cache()

Pytorch 训练时无用的临时变量可能会越来越多，导致 out of memory ，可以使用下面语句来清理这些不需要的变量：torch.cuda.empty_cache() 。
官网 上的解释为：
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. 意思就是PyTorch的缓存分配器会事先分配一些固定的显存，即使实际上tensors并没有使用完这些显存，这些显存也不能被其他应用使用。这个分配过程由第一次CUDA内存访问触发的。而 torch.cuda.empty_cache() 的作用就是释放缓存分配器当前持有的且未占用的缓存显存，以便这些显存可以被其他GPU应用程序中使用，并且通过 nvidia-smi命令可见。注意使用此命令不会释放tensors占用的显存。对于不用的数据变量，Pytorch 可以自动进行回收从而释放相应的显存。
更详细的优化可以查查 优化显存使用 和 显存利用问题。

BUGs

UndefinedMetricWarning: F-score is ill-defined and being set to 0.0 in labels with no predicted samples.'precision', 'predicted', average, warn_for)`

metrics.f1_score(y_test, y_pred, average='weighted')

原因是预测的label都预测成了一个，可能是预测批数据量少了，也可能是模型有问题，从而全部预测成了同一个label。

some labels in y_true don't appear in y_pred. 

[UndefinedMetricWarning: F-score is ill-defined and being set to 0.0 in labels with no predicted]

模型评估

model.eval()的作用是 不启用 Batch Normalization 和 Dropout（实际调用model.train(False) ）。保证 BN 层能够用 全部训练数据 的均值和方差，即测试过程中要保证 BN 层的均值和方差不变。对于 Dropout，model.eval() 是利用到了 所有 网络连接，即不进行随机舍弃神经元。（BN和Dropout固定住，不会取平均，而是用训练好的值，不然test的batch_size过小，很容易就会被BN层影响结果）。

示例

# Set to not-training mode to disable dropout
model.eval()

...

with torch.no_grad():    #禁止梯度的计算

    y_pred = model(x)

    loss = ...

训练时：

# Set back to training mode
model.train(True)  # todo

模型推理和预测

推理示例

with torch.no_grad():  # 禁止梯度的计算
    args = argparse.ArgumentParser().parse_args()
    model_config = ModelConfig(model_config_file)
    vocab = Vocab(vocab_file, vob_num=model_config.vocab_size)
    x = build_features(x, model_config, vocab)

    model = ClassNet()
    model.load_state_dict(torch.load('net_params.pth'))
    model = model.to(args.device)
    model.eval()

    y_pred = model(x, mode="predict")

Note: 如果训练时候是在gpu训练的，预测是在cpu上进行。需要注意几点：
1 模型load时要load成cpu的
model.load_state_dict(torch.load(os.path.join(model_path_name, 'model.bin'),
                                 map_location=None if torch.cuda.is_available() else 'cpu'))
2 模型最好也to到cpu里面
args = argparse.Namespace()
args.device = torch.device("cpu")
model = model.to(args.device)
3 输入数据也to到cpu里面
output = model(input.to('cpu'))

推理加速

非cpu资源独占环境中的torch模型cpu推理加速。

对于某些线上模型在cpu机器上推理时，速度非常慢（但是本地机器非常快，快几百几千倍），原因是本地机器的cpu是独占的，torch会开n个线程，也实际用n个线程，快；但是线上是线程竞争环境，而torch默认使用物理CPU核心数（cat /proc/cpuinfo | grep name | cut -f2 -d: | uniq -c）的线程进行训练，所以会恶化。
解决：改成单线程跑是最安全的。即：

torch_thread_num = int(os.getenv("TORCH_THREAD_NUM", "1"))
torch.set_num_threads(torch_thread_num)

[pytorch中神经网络的多线程数设置：torch.set_num_threads(N)]

一个更完整通用的设置单线程code（基本包含所有可能）

def SetUpPytroch():
    torch_thread_num = int(os.getenv("TORCH_THREAD_NUM", "1"))
    print("torch_thread_num:{}".format(torch_thread_num), file=sys.stderr)
    print("get_num_threads:{}".format(torch.get_num_threads()), file=sys.stderr)
    if torch_thread_num != -1:
        torch.set_num_threads(torch_thread_num)
        torch.set_num_interop_threads(torch_thread_num)
    opt_cudnn = os.getenv("ENABLE_TORCH_OPT_CUDNN", "false")
    if opt_cudnn == "true":
        torch.backends.cudnn.enabled = True
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.deterministic = True
        print("opt_cudnn:{}".format(opt_cudnn), file=sys.stderr)


def SetUpThreadNum():
    thread_num = int(os.getenv("DEFAULT_THREAD_NUM", "1"))
    print("thread_num:{}".format(thread_num), file=sys.stderr)
    os.environ['OMP_NUM_THREADS'] = str(thread_num)
    os.environ['TBB_NUM_THREADS'] = str(thread_num)
    os.environ['OPENBLAS_NUM_THREADS'] = str(thread_num)
    os.environ['MKL_NUM_THREADS'] = str(thread_num)
    os.environ['VECLIB_MAXIMUM_THREADS'] = str(thread_num)
    os.environ['NUMEXPR_NUM_THREADS'] = str(thread_num)
    os.environ['TENSORFLOW_INTER_OP_PARALLELISM'] = str(thread_num)
    os.environ['TENSORFLOW_INTRA_OP_PARALLELISM'] = str(thread_num)

其它可能优化参考

[Performance Tuning Guide — PyTorch Tutorials 2.0.1+cu117 documentation] 

from: -柚子皮-

ref: PyTorch trick 集锦 - 知乎

[9个技巧让你的PyTorch模型训练变得飞快！]

[使用pytorch时，训练集数据太多达到上千万张，Dataloader加载很慢怎么办?]