💻 安装

import Pkg
Pkg.add("Lux")

[!TIP] 如果想在线使用 Lux，可以使用 Google Colab。Colab 已预装了 Julia 运行时环境以及 Lux 和 Reactant！

🚀 基准测试

目前基准测试分散在几个地方：

1. 若要与其他 Julia 包（如 CUDA.jl）进行比较，请查看 Lux.jl/perf。
2. https://enzymead.github.io/Enzyme-JAX/benchmarks/ 展示了 EnzymeJAX（Reactant.jl 的后端）与 JAX 的性能对比。
3. https://enzymead.github.io/Reactant.jl/benchmarks/ 展示了 Reactant.jl 与默认 XLA 和基础 Julia 编译之间的性能对比。

🤸 快速入门

反应物与酶

using Lux, Random, Optimisers, Reactant, Enzyme

rng = Random.default_rng()
Random.seed!(rng, 0)

model = Chain(Dense(128, 256, tanh), Chain(Dense(256, 1, tanh), Dense(1, 10)))

dev = reactant_device()

ps, st = Lux.setup(rng, model) |> dev

x = rand(rng, Float32, 128, 2) |> dev

# 我们需要编译模型才能使用它。
model_forward = @compile model(x, ps, Lux.testmode(st))
model_forward(x, ps, Lux.testmode(st))

# 梯度可以使用Enzyme计算
@jit Enzyme.gradient(Reverse, sum ∘ first ∘ Lux.apply, Const(model), x, ps, Const(st))

# 所有这些都可以通过TrainState API自动化
train_state = Training.TrainState(model, ps, st, Adam(0.001f0))

gs, loss, stats, train_state = Training.single_train_step!(
    AutoEnzyme(), MSELoss(),
    (x, dev(rand(rng, Float32, 10, 2))), train_state
)

原生 Julia 与 Zygote

using Lux, Random, Optimisers, Zygote
# using LuxCUDA, AMDGPU, Metal, oneAPI # GPU支持的可选包

# 种子设置
rng = Random.default_rng()
Random.seed!(rng, 0)

# 构建层
model = Chain(Dense(128, 256, tanh), Chain(Dense(256, 1, tanh), Dense(1, 10)))

# 获取由Lux决定的设备
dev = gpu_device()

# 参数和状态变量
ps, st = Lux.setup(rng, model) |> dev

# 虚拟输入
x = rand(rng, Float32, 128, 2) |> dev

# 运行模型
y, st = Lux.apply(model, x, ps, st)

# 梯度
## 首先构建一个TrainState
train_state = Lux.Training.TrainState(model, ps, st, Adam(0.0001f0))

## 我们可以使用Training.compute_gradients计算梯度
gs, loss, stats, train_state = Lux.Training.compute_gradients(AutoZygote(), MSELoss(),
    (x, dev(rand(rng, Float32, 10, 2))), train_state)

## 优化
train_state = Training.apply_gradients!(train_state, gs) # 或者Training.apply_gradients（没有`!`结尾）

# 这两个步骤可以合并为一次调用
gs, loss, stats, train_state = Training.single_train_step!(AutoZygote(), MSELoss(),
    (x, dev(rand(rng, Float32, 10, 2))), train_state)

📚 示例

请查看examples目录中的独立使用示例。文档中也有按类别整理的示例。

🆘 寻求帮助

如有关于使用的问题，请使用Github Discussions，这样问题和答案都能被索引。如需报告bug，请使用Github Issues，或者更好的是提交一个Pull Request。

🧑‍🔬 引用

如果您在学术工作中发现本库很有用，请引用以下内容：

@software{pal2023lux,
  author    = {Pal, Avik},
  title     = {{Lux: Explicit Parameterization of Deep Neural Networks in Julia}},
  month     = apr,
  year      = 2023,
  note      = {If you use this software, please cite it as below.},
  publisher = {Zenodo},
  version   = {v1.4.2},
  doi       = {10.5281/zenodo.7808903},
  url       = {https://doi.org/10.5281/zenodo.7808903},
  swhid     = {swh:1:dir:1a304ec3243961314a1cc7c1481a31c4386c4a34;origin=https://doi.org/10.5281/zenodo.7808903;visit=swh:1:snp:e2bbe43b14bde47c4ddf7e637eb7fc7bd10db8c7;anchor=swh:1:rel:2c0c0ff927e7bfe8fc8bc43fd553ab392a6eb403;path=/}
}

@thesis{pal2023efficient,
  title     = {{On Efficient Training & Inference of Neural Differential Equations}},
  author    = {Pal, Avik},
  year      = {2023},
  school    = {Massachusetts Institute of Technology}
}

同时，也请给我们的Github仓库点个赞。

🧑‍💻 贡献

本节内容尚不完整。您可以通过完善本节来做出贡献 😜。

💎 格式化（JuliaFormatter）

[!NOTE] 在v2的上游问题解决之前，请将JuliaFormatter固定在v1版本。

using JuliaFormatter
format(".")

🧪 测试

Lux.jl的完整测试耗时较长，以下是测试部分代码的方法。

测试按目录组织，每个目录包含带有@testset块的测试文件。例如，SkipConnection的测试位于test/core_layers/containers_tests.jl中。

运行特定测试文件

运行特定测试最简单的方式是直接激活测试目录并包含该测试文件：

# 从Lux.jl根目录
using Pkg
Pkg.activate("test")

# 运行特定测试文件
include("test/core_layers/containers_tests.jl")

这种方法允许您快速迭代特定测试，而无需运行整个测试套件。

有关执行特定测试组的详细信息，请参阅ParallelTestRunners.jl。

通过CI运行测试组

要通过测试运行程序运行特定测试组，可以将目录名称作为位置参数传递：

julia --project -e 'using Pkg; Pkg.test(test_args=["core_layers"])'

运行所有测试

要运行完整的测试套件：

julia --project -e 'using Pkg; Pkg.test()'

📖 文档

Lux在其文档中创建了许多教程。这可能非常耗时且需要大量计算资源。为了加快构建速度，您可以设置LUX_DOCS_DRAFT_BUILD=true。

LUX_DOCS_DRAFT_BUILD=true julia --threads=auto --startup=no --project=docs docs/make.jl

在编写教程（任何位于examples/下的内容）时，请将其包含在docs/tutorials.jl中。如果教程耗时较长，可以将should_run设置为false。

此外，若要将新页面纳入导航和侧边栏，需要将其添加到docs/src/.vitepress/config.mts中。具体来说，需要根据页面类型将其添加到sidebar和/或nav中。

要使用LiveServer在本地预览文档，请参阅DocumenterVitepress.jl的文档。

Lux.jl 快速上手指南

Lux.jl 是一个用 Julia 语言编写的优雅且高性能的深度学习库，旨在结合 Julia 的编程灵活性与 XLA 的计算性能。

1. 环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：支持 Linux、macOS 或 Windows。
• Julia 版本：建议安装最新稳定版的 Julia (v1.9 或更高版本)。
  • 下载地址：https://julialang.org/downloads/
  • 国内加速：清华大学开源软件镜像站提供 Julia 下载加速 (https://mirrors.tuna.tsinghua.edu.cn/julia-releases/)。
• 前置依赖：无需额外安装 Python 或其他深度学习框架，Lux.jl 是纯 Julia 实现。但若要使用 GPU 加速，需确保已安装对应的 NVIDIA 驱动和 CUDA Toolkit（Lux 会自动处理大部分 CUDA 依赖）。

提示：如果您不想在本地配置环境，可以直接使用 Google Colab，其 Julia 运行时已预装了 Lux 和相关组件。

2. 安装步骤

启动 Julia REPL（交互式命令行），执行以下命令安装 Lux 主包：

import Pkg
Pkg.add("Lux")

安装完成后，建议在代码中通过以下方式加载库：

using Lux, Random

注意：Lux 生态系统包含多个子包（如 LuxLib, LuxCore, MLDataDevices 等），安装主包 Lux 时通常会自动解析并安装必要的核心依赖。

3. 基本使用

以下是一个最简单的示例，展示如何定义一个多层感知机（MLP）模型，初始化参数，并进行一次前向传播。

示例代码

using Lux, Random, ComponentArrays

# 1. 定义模型结构
# 输入层 784 -> 隐藏层 64 (ReLU) -> 输出层 10
model = Chain(
    Dense(784, 64, relu),
    Dense(64, 10)
)

# 2. 初始化随机数生成器和参数
rng = Random.default_rng()
# Lux 需要显式初始化参数和状态
ps, st = Lux.setup(rng, model)

# 将参数转换为组件数组 (可选，但推荐用于优化器兼容)
ps_comp = ComponentArray(ps)

# 3. 准备输入数据 (模拟一个批次的数据)
# 假设输入维度为 (特征数，批次大小)
x = rand(Float32, 784, 32)

# 4. 执行前向传播
# 返回值包括：预测结果 (y) 和更新后的状态 (st_new)
y, st_new = model(x, ps, st)

# 输出结果维度检查
println("输入形状: ", size(x))
println("输出形状: ", size(y))

代码说明

• 模型定义：使用 Chain 串联网络层，Dense 定义全连接层。
• 参数初始化：Lux 采用函数式编程风格，参数 (ps) 和状态 (st) 与模型结构分离，需通过 Lux.setup 初始化。
• 前向传播：调用 model(x, ps, st) 即可得到预测结果，无需像传统 OOP 框架那样调用 .forward() 方法。

现在您已经成功运行了第一个 Lux.jl 模型！接下来您可以结合 Optimisers.jl 进行训练，或使用 LuxCUDA.jl 启用 GPU 加速。