模型训练、评估与推理

作者: MinYao Ni

前面两个章节已经准备好了数据集和模型，本章将讲述如何使用数据训练，评估模型，以及如何使用模型对输入数据进行推理计算。主要包括：

• 模型训练：训练包括多轮迭代（epoch），每轮迭代遍历一次训练数据集，并且每次从中获取一小批（batch）样本，送入模型执行前向计算得到预测值，并计算预测值与真实值之间的损失函数值（loss）。执行梯度反向传播，并根据设置的优化算法（optimizer）更新模型的参数。观察每轮迭代的loss值减小趋势，可判断模型训练效果。

• 模型评估：将测试数据集送入训练好的模型进行评估，得到预测值，计算预测值与真实值之间的损失函数值（loss），并计算评价指标值（metric），便于评估模型效果。

• 模型推理：将数据（样本）送入训练好的模型执行推理，得到推理结果。

还是以手写数字识别任务为例展示训练、评估、推理流程。

模型推理

模型训练主要包含以下的步骤：

1.加载训练数据集，模型定义与初始化。

2.设定义损失函数和评估指标，设置反向传播，选择优化器和设置各项超参

3.初始化优化器状态

4.训练模型，一般包含以下步骤：

• 从dataloader获得一批训练数据

• 对这批训练数据进行推理预测，得到预测结果

• 计算预测结果和真实标签之间的损失值以及评估指标结果

• 通过优化器对损失值进行反向传播，更新模型权重

• 打印模型轮数、批次、损失值、准确率等数据

具体代码如下：

from neurai.grads import BP
import jax.numpy as jnp
from neurai.util.trans import jit
import neurai.nn as nn
from neurai.nn.loss import softmax_cross_entropy
from neurai.opt import apply_updates, adamw
from neurai import datasets

train_dataset = datasets.MNIST(root = './', train=True)
train_loader = datasets.DataLoader(train_dataset, batch_size=64, shuffle=True, drop_last=True)
# 实例化模型
lenet=LeNet()
# 初始化模型权重
model_params = lenet.init(jnp.ones((1,28,28,1)))

# 定义评估指标
def accuracy(predict, target):
  """Calculate accuracy of model output"""
  return jnp.mean(jnp.argmax(predict, axis=1) == jnp.argmax(target, axis=1))

# 设置反向传播， "softmax_cross_entropy"为损失函数
bp = BP(softmax_cross_entropy, accuracy)

# 设置优化器和学习率
optimizer = adamw(learning_rate=0.001)
# 初始化优化器状态
opt_state = optimizer.init(model_params['param'])

@jit # jit加速；以静态图方式运行，速度更快
def train_step(params, batch, opt_state):
  '''
  单步训练
  主要包含模型前向推理、计算loss和指标、获取梯度，更新权重参数等步骤
  '''
  data, label = batch
  data = data.astype(jnp.float32)
  n_label = jnp.asarray(label[:, None] == jnp.arange(10), jnp.float32)
  # 通过反向传播接口，自动完成前向计算，loss和指标计算，
  # 并反向计算梯度，输出梯度，损失值，指标值和当前模型权重
  grads, loss_val, acc_val, params = bp.get_grads(lenet.run,
                                                  params, (data, n_label),
                                                  return_param=True)
  # 按照梯度更新模型权重和优化器状态
  updates, opt_state = optimizer.update(grads['param'], opt_state,
                                        params['param'])
  newparam = apply_updates(params['param'], updates)
  params['param'] = newparam

  return params, opt_state, loss_val, acc_val

epochs = 1
for epoch in range(epochs):
  for  batch_id, batch in enumerate(train_loader):
    model_params, opt_state, loss_val, acc_val = train_step(model_params, batch, opt_state)
    if batch_id%100 == 0:
      print("epoch: {}, batch_id: {}, loss: {}, accuracy: {}".format(
                                                                      epoch,
                                                                      batch_id,
                                                                      loss_val,
                                                                      acc_val
                                                                    ))

epoch: 0, batch_id: 0, loss: 201.3196258544922, accuracy: 0.078125
epoch: 0, batch_id: 100, loss: 1.7603965997695923, accuracy: 0.828125
epoch: 0, batch_id: 200, loss: 0.4224163293838501, accuracy: 0.90625
epoch: 0, batch_id: 300, loss: 1.4008324146270752, accuracy: 0.90625
epoch: 0, batch_id: 400, loss: 0.8987967371940613, accuracy: 0.921875
epoch: 0, batch_id: 500, loss: 0.3118555247783661, accuracy: 0.90625
epoch: 0, batch_id: 600, loss: 0.18040047585964203, accuracy: 0.9375
epoch: 0, batch_id: 700, loss: 0.45215797424316406, accuracy: 0.953125
epoch: 0, batch_id: 800, loss: 0.44177737832069397, accuracy: 0.9375
epoch: 0, batch_id: 900, loss: 0.1789383739233017, accuracy: 0.9375

模型评估

在这个阶段，将对训练得到的模型进行效果评估，一般会使用同分布但未参与训练的数据测试模型效果。模型评估和训练有一定的差异，主要流程为：

1.批量加载测试数据集

2.将批数据输入模型，得到推理结果

3.计算损失与评估指标

和训练的主要区别在于：不需要在这个阶段更新模型权重，也就不用定义优化器、计算梯度和反向传播等。

以训练后LeNet模型的评估为例，代码如下：

test_dataset = datasets.MNIST(root = './', train=False)
test_loader = datasets.DataLoader(test_dataset, batch_size=64, shuffle=True, drop_last=True)

@jit
def test_step(params, batch):
  '''
  单步验证
  主要包含模型推理，以及损失和指标计算
  '''
  data, label = batch
  data = data.astype(jnp.float32)
  n_label = jnp.array(label[:, None] == jnp.arange(10), jnp.float32)
  # 模型仅推理，并计算loss和准确率指标
  loss_val, acc_val = bp.predict(lenet.run,
                                params, (data, n_label))
  return loss_val, acc_val

total_acc = 0.0
total_loss = 0.0
total_num = 0
for batch_id, batch in enumerate(test_loader):
  batch_size = batch[0].shape[0]
  loss_val, acc_val = test_step(model_params, batch)
  # 统计准确率
  total_acc = total_acc + acc_val * batch_size
  # 统计损失值
  total_loss = total_loss + loss_val * batch_size
  total_num = batch_size + total_num
acc = total_acc / total_num
loss = total_loss / total_num
print("test loss: {}, test accuracy: {}".format(loss, acc))

test loss: 0.27429306507110596, test accuracy: 0.9448117017745972

模型推理

和模型的训练评估相比，模型推理流程更为简单，可以认为模型推理就是模型实际使用的过程。具体步骤如下：

1.将需要执行计算的数据，预处理为模型输入需要的格式

2.将数据输入模型得到预测结果

3.将模型输出结果进行后处理得到最终结果

test_loader = datasets.DataLoader(test_dataset, batch_size=1, shuffle=True, drop_last=True)
for data in test_loader:
  img, label = data
  img = img.astype(jnp.float32)
  # 模型预测
  predict = lenet.run(model_params, img)
  # 预测结果后处理
  predict_label = jnp.argmax(predict[0])
  print("true label: {}, predict label: {}".format(label, predict_label))
  cv2.imwrite("./mnist-0.png", np.asarray(img[0,...,0], np.int8))
  break

true label: [0], predict label: 0