Flux.jl 深度学习教程：使用多层感知机(MLP)实现MNIST手写数字分类-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00273/article/details/148523841

Flux.jl 深度学习教程：使用多层感知机(MLP)实现MNIST手写数字分类

Flux.jl Relax! Flux is the ML library that doesn't make you tensor 项目地址: https://gitcode.com/gh_mirrors/fl/Flux.jl

前言

多层感知机(MLP)是最基础的神经网络结构之一，广泛应用于各种分类和回归任务。本文将详细介绍如何使用Flux.jl框架构建一个简单的MLP模型，并在经典的MNIST手写数字数据集上进行训练和评估。

环境准备

首先需要安装必要的Julia包：

using Flux, Statistics
using Flux: DataLoader
using Flux: onehotbatch, onecold, logitcrossentropy
using MLDatasets: MNIST
using MLUtils

对于有GPU的用户，可以添加CUDA支持以加速训练过程。

模型参数配置

我们定义一个结构体来存储训练参数：

Base.@kwdef mutable struct Args
    rate::Float64 = 3e-4    # 学习率
    batchsize::Int = 1024   # 批量大小
    epochs::Int = 10        # 训练轮数
    usegpu::Bool = true     # 是否使用GPU
end

这些参数包括学习率、批量大小、训练轮数以及是否使用GPU的选项。合理的参数设置对模型训练效果至关重要。

数据准备

数据预处理是机器学习流程中的关键步骤。我们定义getdata函数来加载和准备MNIST数据集：

function getdata(args)
    # 加载数据集
    xtrain, ytrain = MNIST(:train)[:]
    xtest, ytest = MNIST(:test)[:]
    
    # 将图像数据展平为向量
    xtrain = Flux.flatten(xtrain)
    xtest = Flux.flatten(xtest)

    # 对标签进行one-hot编码
    ytrain, ytest = onehotbatch(ytrain, 0:9), onehotbatch(ytest, 0:9)

    # 创建数据加载器
    train_loader = DataLoader((xtrain, ytrain), batchsize=args.batchsize, shuffle=true)
    test_loader = DataLoader((xtest, ytest), batchsize=args.batchsize)

    return train_loader, test_loader
end

这个函数完成了以下工作：

加载MNIST训练集和测试集
将28×28的图像展平为784维向量
将数字标签转换为one-hot编码
创建批量数据加载器，训练数据会随机打乱

模型构建

我们构建一个简单的三层MLP模型：

function build_model(; imgsize=(28,28,1), nclasses=10)
    return Chain(
        Dense(prod(imgsize) => 32, relu),  # 输入层到隐藏层
        Dense(32 => nclasses)              # 隐藏层到输出层
    )
end

模型结构说明：

输入层：784个神经元（对应展平后的28×28图像）
隐藏层：32个神经元，使用ReLU激活函数
输出层：10个神经元（对应0-9十个数字类别）

损失函数与评估指标

定义损失函数来评估模型性能：

function loss_all(dataloader, model)
    l = 0f0
    n = 0
    for (x, y) in dataloader
        l += logitcrossentropy(model(x), y, agg=sum)
        n += MLUtils.numobs(x)
    end
    return l / n
end

这里使用logit交叉熵损失函数，它适用于多分类问题，并且对模型的原始输出（logits）进行操作。

定义准确率计算函数：

function accuracy(dataloader, model)
    acc = 0
    n = 0
    for (x, y) in dataloader
        acc += sum(onecold(cpu(model(x))) .== onecold(cpu(y)))
        n += MLUtils.numobs(x)
    end
    return acc / n
end

准确率是分类任务中最直观的评估指标，表示模型预测正确的比例。

模型训练

完整的训练流程如下：

function train(; kws...)
    # 初始化参数
    args = Args(; kws...)
    device = args.usegpu ? Flux.get_device() : Flux.get_device("CPU")
    
    # 加载数据
    train_loader, test_loader = getdata(args)

    # 构建模型并移至设备
    model = build_model() |> device

    # 定义损失函数
    loss(model, x, y) = logitcrossentropy(model(x), y)
    
    # 设置优化器
    opt_state = Flux.setup(Adam(args.rate), model)
    
    # 训练循环
    for epoch in 1:args.epochs
        @info "Epoch $epoch"
        for d in train_loader
            x, y = d |> device
            g = gradient(m -> loss(m, x, y), model)[1]
            Flux.update!(opt_state, model, g)
        end
        @show accuracy(train_loader, model)
        @show accuracy(test_loader, model)
    end
end

# 开始训练
train()

训练过程包括：