Flux | Julia 机器学习框架入门实战

Flux 是一个机器学习框架。它最大的特点是配合 Julia 的语法，用编写可微分算法的思想设计机器学习算法，通过更为抽象的数学形式来代替嵌入操作。

Flux 与机器学习

Python 没有特地为科学计算或者说机器学习有过倾斜。Julia 却从一开始就是为了科学计算而生，在语法层面、编译器角度都为科学计算打好了坚实的基础，非常适合表达机器学习的程序。

Flux希望基于 Julia 语法，重新思考机器学习算法的编写。Flux 的切入点是目前的机器学习算法本质上其实是可微分算法，这和传统的离散数据结构是截然不同的。说到底，无非是在离散结构机器上模拟计算连续的算法。因此如果能更加数学化地表达可微分算法，也就可以更为高效简洁地编写机器学习算法。 Flux 的开发者还发现，求微分的过程，本质上是符号转换，这其实很大程度上是编写编译器的时候需要解决的问题。因此如果能拓展编译器的实现，也就可以在语言层面解决问题。幸运地是，Julia 提供了这种可能，因此 Flux 的开发者开发出了 Zygote(21世纪的自动微分)，作为 Flux 的基础。凭借 Julia 强大的基础设施，不再需要数十万行的 C 代码，Flux 仅仅只有数千行的 Julia 代码（当然，成熟程度另当别论）。

有关 Flux 设计的问题可以参考 Julia 官网博客文章 Building a Language and Compiler for Machine Learning 或者在 arxiv 上阅读论文 Fashionable Modelling with Flux

机器学习介绍

去年，公众号已经推送了一篇通过纯粹的矩阵运算，没有优化地实现全连接的三层神经网络，并在最经典的 MNIST 集上进行训练，通过简单的多轮循环训练，达到了 MNIST 官网给出的水准（P.S. 训练速度相对于使用框架而言是非常慢的）。在那篇文章中，对于机器学习已经有比较粗浅的介绍，因此就不重复了。简而言之，机器学习还是相当的流水化，很有几分八股文的味道。步骤是固定的，无非是获取数据、洗数据，理解数据并搭建模型，测试集／训练集划分，训练，优化。

获取 MNIST 数据

Flux 已经内置了 MNIST 的数据加载器（ Flux.Data.MNIST)，可以直接下载 MNIST 经典的 60000 测试图片，相当方便。

1. using Flux.Data.MNIST

2. imgs = MNIST.images()

默认上，这是训练集。可以通过传入 :test (元编程技术)获取 10000 的测试集图片。

imgs = MNIST.images(:test)

得到的数据是一个嵌套的数组，为了方便操作和性能考虑，可以重新将其变形为一维度的向量。可以直接使用 vec 或者使用 reshape：

1. # 两种写法

2. X = hcat(float.(vec.(imgs))...)

3. X = hcat(float.(reshape.(imgs, :))...)

hcat 把 [A;B;C;...] 转变成 [A B C...]，也就是从一维的竖直向量转变为水平方向。 接下来，获取每张图片对应的数字。

labels = MNIST.labels()# 同样，传入 `:test` 可以获取测试集对应标签labels = MNIST.labels(:test)

返回的 0 - 9 数字，为了提高效率，可以使用 one-hot 技术（独热编码）编码向量。举个例子：

1. julia> using Flux: onehotbatch

2. julia> onehotbatch([:b, :a, :b], [:a, :b, :c])

3. 3×3 Flux.OneHotMatrix:

4. false true false

5.  true false true

6. false false false

可以这样理解，对于每一个特征，如果它有 m 个可能值，那么经过 one-hot 编码后，就变成了 m 个二元特征。并且，这些特征互斥，每次只有一个激活。因此，数据会变成稀疏的。这样就解决了分类器不好处理属性数据的问题，也一定程度上扩充了特征。 因此，我们可以得到这样编写代码：

using Flux: onehotbatchY = onehotbatch(labels, 0:9)

好了，将上面的代码封装一下，最终写成：

1. # 获取数据集

2. function getdataset(;train=true)

3.    train_or_test = ifelse(train, :train, :test)

4.    imgs = MNIST.images(train_or_test)

5.    X = hcat(float.(vec.(imgs))...)

6.    labels = MNIST.labels(train_or_test)

7.    Y = onehotbatch(labels, 0:9)

8.    return X, Y

9. end

搭建模型

MNIST 数据当然不用洗了，可以开始搭建模型了，Flux 内置了一些层次模型，比如卷积、池化等等，我们这次用最普通的三层全链接神经网络。三层神经网络模型的中间隐藏层规模是可以随意的，根据矩阵运算性质可以屏蔽调该参数。

注意，28^2 和10 是不能变的。前者是 MNIST 图片是 28 × 28 大小的图片，后者是 0 - 9 这10 个数字。

function define_model(;hidden)    mlp = Chain(Dense(28^2, hidden, relu),                Dense(hidden, hidden, relu),                Dense(hidden, 10),                softmax)    return mlpend

relu 是线性整流函数（Rectified Linear Unit, ReLU），作为激活函数。Flux 还提供了其它的激活函数，比如经典的 Sigmoid 激活函数，还有 Kaggle 中被首先提出并使用的泄漏随机线性整流函数 leakyrelu，等等。

至于 softmax，可以理解为归一化的意思。识别0-9这10个手写数字，若最后一层的输出为[0,1,0, 0, 0, 0, 0, 0, 0, 0]，则表明我们网络的识别结果为数字1。 

划分训练集和测试集

训练过程要量化训练效果，就需要使用测试集来检验。因此我们将训练和测试的比例随机分成 9:1。

1. function split_dataset_random(X, Y)

2.    divide_ratio = 0.9 # 划分比例

3.    shuffled_indices = shuffle(1:size(Y)[2]) # 随机打乱

4.    divide_idx = round(Int, 0.9 * length(shuffled_indices)) # 计算切分点

6.    train_indices = shuffled_indices[1:divide_idx] # 1 : 切分点 作为训练集

7.    val_indices = shuffled_indices[divide_idx:end] # 切分点 : end 作为测试集

8.    train_X = X[:,train_indices] # 数据

9.    train_Y = Y[:,train_indices] # 对应标签

10.    val_X = X[:,val_indices]

11.    val_Y = Y[:,val_indices]

12.    return train_X, train_Y, val_X, val_Y

13. end

如果觉得太麻烦，那么直接重复放大数据集也是可以的 repeated((X,Y),200)。

训练模型

终于到了最核心的地方，在 train 函数中，需要结合上面的准备工作。 首先，准备训练数据，接着划分训练集和测试集。

X, Y = prepare_dataset(train=true)train_X, train_Y, val_X, val_Y = split_dataset_random(X, Y)

定义模型、损失函数和准确率函数：

1. model = define_model(hidden = hidden)

3. loss(x,y)= crossentropy(model(x),y)

4. accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))    

crossentropy，就是交叉熵，可以衡量两者的相似度。 mean 是取平均值，值得一提的是 onecold，这和前的 onehot 显然是一个相反的过程，onehot 把特征进行编码，onecold 自然是把解码成原来的特征。 .== 便可以广播到所有的元素，得到一个布尔值的向量，经过平均值 mean 后便得到了准确度。

接下来，将数据以 64 个一组进行划分：

batchsize = 64train_dataset = [(train_X[:,batch] ,train_Y[:,batch]) for batch in partition(1:size(train_Y)[2],batchsize)])val_dataset = [(val_X[:,batch] ,val_Y[:,batch]) for batch in partition(1:size(val_Y)[2],batchsize)]

定义优化器，这里使用的 Flux 内置的 ADAM 算法，这是一个随机优化算法，可以替代传统的随机梯度下降的一阶优化算法，能够基于训练数据的迭代更新神经网络的权重。

1. optimizer = ADAM(params(model))

Flux 中，ADAM 的原型是这样的：

ADAM(params, η = 0.001; β1 = 0.9, β2 = 0.999, ϵ = 1e-08, decay = 0)

在 ADAM 算法的论文中，也给出了在 MNIST 集上的运行效果：

关于 ADAM 算法的更多信息，可以到 https://arxiv.org/abs/1412.6980 查看相应的论文 Adam: A Method for Stochastic Optimization。

损失函数有了，优化方法也有了，那么接下来就是训练了。为了获得比较好的效果，我们依然采取循环训练的手段，将每一轮得到的模型，重新进入到下一轮的训练当中，而这仅仅只需要一个 @epochs。 为了能切实地感受到训练的过程，不妨加上一个回调函数，每一轮训练之后，输出总体的损失值和准确度。借此我们可以看到哪里过拟合，指导后续的调优。除此之外，还能用来保存模型。 

我们使用闭包来编写回调函数，这样可以封装一个私有的计数变量：

1. function make_callback()

2.    callback_count = 0

4.    function callback()

5.        callback_count = 1

6.        if callback_count == length(train_dataset)

7.            println('action for each epoch')

8.            total_loss = 0

9.            total_acc = 0

10.            for (vx, vy) in val_dataset

11.                total_loss = loss(vx, vy)

12.                total_acc = accuracy(vx, vy)

13.            end

14.            total_loss /= length(val_dataset)

15.            total_acc /= length(val_dataset)

16.            @show total_loss, total_acc

17.            callback_count = 0

18.        end

19.    end

20. end

22. eval_callback = make_callback()

之前每 64 划分一组， callback_count 就是用来计数是否达到了最后一组（即一轮结束）。这样，每一轮都会输出类似这样的结果：

action for each epoch(total_loss, total_acc) = (0.09762976779695344 (tracked), 0.97302092379505)

保存模型

将每一轮模型保存下来，这样即使中断也会有一个模型，还能借此实现断点续训的功能。只需要在 callback 后面加上这样的代码：

1. using BSON: @load, @save

2. pretrained = model

3. @save 'pretrained.bson' pretrained

OK，开始训练：

@epochs epochs Flux.train!(loss, train_dataset, optimizer, cb = eval_callback)

结束之后，保存下最终的模型：

1. pretrained = model

2. weights = Tracker.data.(params(pretrained))

3. @save 'pretrained.bson' pretrained

4. @save 'weights.bson' weights

这里采取两种保存模型的方法，一种是保存整个模型，另一种只是保留结点的权重。相对而言，前者所需的空间更大一点，但是后期使用的时候不再需要创建模型，再进行赋值。（可以只选择其中一种来保存模型）

# 加载模型并测试效果function predict()    println('Start to evaluate testset')    println('loading pretrained model')    @load 'pretrained.bson' pretrained    model = pretrained    accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))    println('prepare dataset')    X, Y = prepare_dataset(train=false)    X = X    Y = Y    @show accuracy(X, Y)    println('Done')end# 加载权重并测试效果function predict2()    println('Start to evaluate testset')    println('loading pretrained model')    @load 'weights.bson' weights    model = define_model(hidden=100)  # 创建模型    Flux.loadparams!(model, weights)    model = model    accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))    println('prepare dataset')    X, Y = prepare_dataset(train=false)    X = X    Y = Y    @show accuracy(X, Y)    println('Done')end

效果

最终代码

如果需要使用显卡训练，可以使用 CUDA 配合 CuArrays，如果没有 CUDA，注释掉即可。以下代码支持在 GPU 下训练（和上面的代码稍微有一点增加，例如 |>gpu，把数据重定向到 GPU） ：

1. using Flux, Flux.Data.MNIST, Statistics

2. using Flux: onehotbatch, onecold, crossentropy, @epochs

3. using Base.Iterators: partition

4. using BSON: @load, @save

5. # using CuArrays

6. using Random

8. function prepare_dataset(;train=true)

9.    train_or_test = ifelse(train, :train, :test)

10.    imgs = MNIST.images(train_or_test)

11.    X = hcat(float.(vec.(imgs))...)

12.    labels = MNIST.labels(train_or_test)

13.    Y = onehotbatch(labels, 0:9)

14.    return X, Y

15. end

17. function define_model(;hidden=100)

18.    mlp = Chain(Dense(28^2, hidden, relu),

19.                Dense(hidden, hidden, relu),

20.                Dense(hidden, 10),

21.                softmax)

22.    return mlp

23. end

25. function split_dataset_random(X, Y)

26.    divide_ratio = 0.9 # 划分比例

27.    shuffled_indices = shuffle(1:size(Y)[2]) # 随机打乱

28.    divide_idx = round(Int, 0.9 * length(shuffled_indices)) # 计算切分点

30.    train_indices = shuffled_indices[1:divide_idx] # 1 : 切分点 作为训练集

31.    val_indices = shuffled_indices[divide_idx:end] # 切分点 : end 作为测试集

32.    train_X = X[:,train_indices] # 数据

33.    train_Y = Y[:,train_indices] # 对应标签

34.    val_X = X[:,val_indices]

35.    val_Y = Y[:,val_indices]

36.    return train_X, train_Y, val_X, val_Y

37. end

39. function train(;epochs = 10, hidden = 100)

40.    println('Start to train')

41.    X, Y = prepare_dataset(train=true)

42.    train_X, train_Y, val_X,val_Y = split_dataset_random(X, Y)

43.    model = define_model(hidden = hidden) |> gpu

45.    loss(x,y)= crossentropy(model(x),y)

46.    accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))

48.    batchsize = 64

49.    train_dataset = gpu.([(train_X[:,batch] ,train_Y[:,batch]) for batch in partition(1:size(train_Y)[2],batchsize)])

50.    val_dataset = gpu.([(val_X[:,batch] ,val_Y[:,batch]) for batch in partition(1:size(val_Y)[2],batchsize)])

52.    function make_callback()

53.        callback_count = 0

55.        function callback()

56.            callback_count = 1

57.            if callback_count == length(train_dataset)

58.                println('action for each epoch')

59.                total_loss = 0

60.                total_acc = 0

61.                for (vx, vy) in val_dataset

62.                    total_loss = loss(vx, vy)

63.                    total_acc = accuracy(vx, vy)

64.                end

65.                total_loss /= length(val_dataset)

66.                total_acc /= length(val_dataset)

67.                @show total_loss, total_acc

68.                pretrained = model |> cpu

69.                @save 'pretrained.bson' pretrained

70.                callback_count = 0

71.            end

72.        end

74.        return callback

75.    end

77.    eval_callback = make_callback()

78.    optimizer = ADAM(params(model))

80.    @epochs epochs Flux.train!(loss, train_dataset, optimizer, cb = eval_callback)

82.    pretrained = model |> cpu

83.    weights = Tracker.data.(params(pretrained))

84.    @save 'pretrained.bson' pretrained

85.    @save 'weights.bson' weights

86.    println('Finished to train')

87. end

89. function predict()

90.    println('Start to evaluate testset')

91.    println('loading pretrained model')

92.    @load 'pretrained.bson' pretrained

93.    model = pretrained |> gpu

94.    accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))

95.    println('prepare dataset')

96.    X, Y = prepare_dataset(train=false)

97.    X = X |> gpu

98.    Y = Y |> gpu

99.    @show accuracy(X, Y)

100.    println('Done')

101. end

103. function predict2()

104.    println('Start to evaluate testset')

105.    println('loading pretrained model')

106.    @load 'weights.bson' weights

107.    model = define_model(hidden=100)

108.    Flux.loadparams!(model, weights)

109.    model = model |> gpu

110.    accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))

111.    println('prepare dataset')

112.    X, Y = prepare_dataset(train=false)

113.    X = X |> gpu

114.    Y = Y |> gpu

115.    @show accuracy(X, Y)

116.    println('Done')

117. end

119. function main()

120.    train()

121.    predict()

122.    predict2()

123. end

125. main()

手写数字实验

接下来看看效果怎么样，用鼠标绘制了28×28的手写数字。另外单独写一段代码调用一下训练好的模型：

using Images, ImageViewusing Fluxusing Flux: onecoldusing BSON: @load, @save# 加载模型@load 'pretrained.bson' pretrainedmodel = pretrainedlabel = 4img = load('num_imgs/$label.png')# 预处理图片test_img = convert(Array{Float64}, 1 * Gray.(img) .> 0.5);X = float.(vec(test_img));onecold(model(X))onecold(model(X)) == label

 

最后

相对于之前纯粹的未经性能优化的原始三层神经网络，使用 Flux 的速度快得多。代码上，可以看到 Flux 在 Julia 的加持下还是很简洁地完成了搭建模型到训练的整个工作。最终也获得了 97.75 这样相当不错的成绩（可以进一步根据 LeCun 官网上的参数进行调整）。