模型训练

在进行模型训练时，一般分为四个步骤。

构建数据集，参见数据集 Dataset。
定义神经网络，参见网络构建。
定义超参、损失函数及优化器。
训练网络，参见训练及保存模型。

准备数据集和神经网络

参考前面章节数据集 Dataset和网络构建。给出如下加载数据集和定义神经网络的代码

from CangjieTB import macros.*
from CangjieTB import nn.*
from CangjieTB import nn.layers.*
from CangjieTB import nn.optim.*
from CangjieTB import nn.loss.*
from CangjieTB import ops.*
from CangjieTB import common.*
from CangjieTB import dataset.*

// 定义神经网络
@OptDifferentiable
struct LeNetModel {
    let conv1: Conv2d
    let conv2: Conv2d
    let pool: MaxPool2d
    let dense1: Dense
    let dense2: Dense
    let dense3: Dense
    init() {
        conv1 = Conv2d(1, 6, Array<Int64>([5, 5]), pad_mode: "valid")
        conv2 = Conv2d(6, 16, Array<Int64>([5, 5]), pad_mode: "valid")
        pool = MaxPool2d(kernel_size: Array<Int64>([2, 2]), stride: Array<Int64>([2, 2]))
        dense1 = Dense(400, 120, RandomNormalInitializer(sigma: 0.02))
        dense2 = Dense(120, 84, RandomNormalInitializer(sigma: 0.02))
        dense3 = Dense(84, 10, RandomNormalInitializer(sigma: 0.02))
    }
@Differentiable
    operator func ()(input: Tensor): Tensor {
        input |> this.conv1 |> relu |> this.pool
            |> this.conv2 |> relu |> this.pool
            |> flatten
            |> this.dense1 |> relu
            |> this.dense2 |> relu
            |> this.dense3
    }
}

func createLenetDataset() {
    // 加载数据集
    let dataPath: String = "./data/mnist/test"
    var mnistDs = MnistDataset(dataPath, epoch: 2)
    // 将图像数据映射到[0, 1]的范围
    var rescale = rescale(1.0 / 255.0, 0.0)
    mnistDs.datasetMap([rescale], "image")
    return mnistDs
}

定义超参、损失函数及优化器

定义超参

超参是可以调整的参数，可以控制模型训练优化的过程，不同的超参数值可能会影响模型训练和收敛速度。

一般会定义以下用于训练的超参：

训练轮次（epoch）：训练时遍历数据集的次数。
批次大小（batch size）：数据集进行分批读取训练，设定每个批次数据的大小。
学习率（learning rate）：如果学习率偏小，会导致收敛的速度变慢，如果学习率偏大则可能会导致训练不收敛等不可预测的结果。

为数据集设置训练轮次和批次大小，定义学习率

let epoch: Int32 = 5
let batchSize: Int32 = 32
mnistDs.batch(batchSize, true)
mnistDs.repeat(epoch)

let learningRate: Float32 = 1.0e-2

设置学习率见定义优化器。

定义损失函数

使用损失函数需要导入 nn 包。

from CangjieTB import nn.loss.*

损失函数用来评价模型的预测值和真实值不一样的程度。给定输出值和目标值，计算损失值，使用方法如下所示：

// 创建损失函数对象
let lossFn = SoftmaxCrossEntropyWithLogits(sparse: true)
// 创建计算值与目标值
let netOutput = Tensor(Array<Float32>([0.1, 0.3, 0.5, 0.1]), shape: Array<Int64>([1, 4]))
let label = Tensor(Array<Int32>([3]), shape: Array<Int64>([1]))
// 调用损失函数对象的操作符重载方法，获取 loss 值
let loss = lossFn(netOutput, label)
print("loss", loss)

输出为：

loss
Tensor(shape=[], dtype=Float32, value= 1.55037284e+00)

除了 SoftmaxCrossEntropyWithLogits，仓颉 TensorBoost 也支持用户自定义损失函数。与自定义神经网络结构类似，自定义损失函数使用 struct 数据类型，其计算逻辑在 operator() 运算符重载函数中实现，详见构建神经网络。

定义优化器

使用优化器需要导入 nn 包。

from CangjieTB import nn.*

优化器用于计算和更新梯度，构建一个优化器对象，该对象可以基于计算得到的梯度对原模型进行参数更新。设置优化器的参数选项，比如学习率、权重衰减等等。 class MomentumOptimizer<T0> 中泛型 T0 为神经网络的类型，如 LeNetModel 或 MyLayer 等。初始化时必须显式指明泛型类型，代码如下：

let lr: Float32 = 0.01
let momentum: Float32 = 0.9
let net = LeNetModel()
var optim: MomentumOptimizer<LeNetModel> = MomentumOptimizer<LeNetModel>(net, lr, momentum)

其中，学习率可以是 Float32 类型的常数，也可以是 Tensor 类型的可变动态学习率，代码如下：

var lrArrary: Array<Float32> = Array<Float32>([0.01, 0.01, 0.005, 0.002, 0.001])
var learningRate = Tensor(lrArrary, shape: [lrArrary.size])
var net = LeNetModel()
var optim = MomentumOptimizer<LeNetModel>(net, learningRate, 0.9, weightDecay: 0.00004, lossScale: 1024.0)

训练网络

训练网络的过程可以分成两步，第一步是进行前向和反向计算分别得到模型的 loss 值和梯度；第二步是利用梯度更新模型参数。

计算 loss 值和梯度

运用仓颉 TensorBoost 提供的 @AdjointOf 接口，可以获取到网络反向计算的梯度，代码如下：

// 定义可微的训练函数
@Differentiable[except: [lossFn, input, label]]
func train(net: LeNetModel, lossFn: SoftmaxCrossEntropyWithLogits, input: Tensor, label: Tensor): Tensor
{
    var output = net(input)
    var lossTensor = lossFn(output, label)
    return lossTensor
}

根据训练函数，获取输入的梯度：

// 创建 train 函数的输入
var input = parameter(zerosTensor(Array<Int64>([32, 1, 32, 32]), dtype: FLOAT32), "data")
var label = parameter(zerosTensor(Array<Int64>([32]), dtype: INT32), "label")
var lossFn = SoftmaxCrossEntropyWithLogits(sparse: true)
// 获取伴随函数
let adj = @AdjointOf(train)
// 获取反向传播函数
var (loss, bpFunc) = adj(net, lossFn, input, label)
// 获取梯度
var gradient: LeNetModel = bpFunc(Tensor(Float32(1.0)))

更新参数

当完成优化器对象的初始化，并通过反向传播函数得到梯度对象之后，即可调用优化器的 update 方法，将获取到的梯度传递给优化器，代码如下：

// 更新梯度
optim.update(gradient)

此处的 gradient 的类型需要与 optim 对象的泛型一致。

执行训练

基于前面介绍的梯度计算和参数更新方法，可以写出如下训练代码:

var datasetLenetTest = createLenetDataset()
for (i in 1..(datasetLenetTest.getDatasetSize() + 1)) {
    datasetLenetTest.getNext([input, label])
    var (loss, bpFunc) = adj(net, lossFn, input, label)
    var gradient: LeNetModel = bpFunc(Tensor(Float32(1.0)))
    optim.update(gradient)
    let lossValue = loss.evaluate().toScalar<Float32>()   // 这里调用 `evaluate` 确保静态图模式下能够正确的执行计算
    if (i % 100 == 0) {
        print("step: ${i}, loss is: ${lossValue}\n")
    }
}

这就是一个神经网络模型的完整训练过程了。执行如上代码，一般可以看到模型 loss 值逐步减小，代表模型收敛。如果遇到模型不收敛的情况，可以尝试调整超参。