The 🔥 Deep Learning Framework
] add Lux
using Lux, Random, Optimisers, Zygote
# Seeding
rng = Random.default_rng()
Random.seed!(rng, 0)
# Construct the layer
model = Chain(BatchNorm(128), Dense(128, 256, tanh), BatchNorm(256),
Chain(Dense(256, 1, tanh),Dense(1, 10)))
# Parameter and State Variables
ps, st = Lux.setup(rng, model) .|> gpu
# Dummy Input
x = rand(rng, Float32, 128, 2) |> gpu
# Run the model
y, st = Lux.apply(model, x, ps, st)
# Gradients
gs = gradient(p -> sum(Lux.apply(model, x, p, st)[1]), ps)[1]
# Optimization
st_opt = Optimisers.setup(Optimisers.ADAM(0.0001), ps)
st_opt, ps = Optimisers.update(st_opt, ps, gs)
Look in the examples directory for self-contained usage examples. The documentation has examples sorted into proper categories.
Checkout our Ecosystem page for more details.
For usage related questions, please use Github Discussions or JuliaLang Discourse (machine learning domain) which allows questions and answers to be indexed. To report bugs use github issues or even better send in a pull request.
- Flux.jl -- We share most of the backend infrastructure with Flux (Roadmap hints towards making Flux explicit-parameter first)
- Knet.jl -- One of the mature and OG Julia Deep Learning Frameworks
- SimpleChains.jl -- Extremely Efficient for Small Neural Networks on CPU
- Avalon.jl -- Uses tracing based AD Yota.jl
If you found this library to be useful in academic work, then please cite:
@software{pal2022lux,
author = {Pal, Avik},
title = {{L}ux: Explicit Parameterization of Deep Neural Networks in Julia},
year = {2022},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://github.com/avik-pal/Lux.jl/}}
}
Also consider starring our github repo