# Best libraries for experimenting with small scale models?
**URL:** https://discourse.julialang.org/t/best-libraries-for-experimenting-with-small-scale-models/128106
**Category:** Machine Learning
**Tags:** question
**Created:** [April 15, 2025, 6:28pm UTC](https://discourse.julialang.org/t/best-libraries-for-experimenting-with-small-scale-models/128106 "2025-04-15T18:28:17Z")
**Posts on this page:** 8
**Page:** 1
### Author:  [@stur86](https://discourse.julialang.org/u/stur86)
#### Post date: [April 15, 2025, 6:28pm UTC](https://discourse.julialang.org/t/best-libraries-for-experimenting-with-small-scale-models/128106/1 "2025-04-15T18:28:17Z")
I’m looking for tools to do ML experiments in Julia. Specifically, I’d like trying to design some small-ish models (hundreds to thousands of parameters) and experiment with different approaches that aren’t just differently ordered chains of the usual NN layers. Basically what I’d like is a framework that provides me with:
- automatic differentiation
- standard optimization algorithms
- loss functions
and very little more, leaving otherwise almost full freedom to write your model as any function. Also, running on CPU is perfectly fine. Is there something that fits these requirements? Lux.jl seemed like a possible choice but I haven’t dug deep into it.
---
### Author:  [@jbrea](https://discourse.julialang.org/u/jbrea)
#### Post date: [April 16, 2025, 8:50am UTC](https://discourse.julialang.org/t/best-libraries-for-experimenting-with-small-scale-models/128106/2 "2025-04-16T08:50:37Z")
Lux.jl is certainly a good option. Depending on what you need exactly, you may also be fine with just using DifferentiationInterface.jl (for autodiff), any optimization package (Optimisers.jl, Optim.jl, Optimization.jl), and custom loss functions or the losses in Flux.jl, Lux.jl or LossFunctions.jl
---
### Author:  [@mcabbott](https://discourse.julialang.org/u/mcabbott)
#### Post date: [April 21, 2025, 8:13pm UTC](https://discourse.julialang.org/t/best-libraries-for-experimenting-with-small-scale-models/128106/3 "2025-04-21T20:13:40Z")
My completely unbiased recommendation is that you try Flux… or just Optimisers.jl + Zygote.jl (or Enzyme.jl).
Flux regards any callable struct containing parameters as a model. The readme gives this example, in which the anonymous function `x -> ...` captures 3 arrays:
```julia
model = let
w, b, v = (randn(Float32, 23) for _ in 1:3) # parameters
x -> sum(v .* tanh.(w*x .+ b)) # callable
end
typeof(model) # var"#52#54"{Vector{Float32}, Vector{Float32}, Vector{Float32}}
```
If you take a derivative with respect to `model` using Zygote.jl (or Enzyme.jl), you get another struct which contains matching fields:
```julia
data = [(x, 2x-x^3) for x in -2:0.1f0:2];
grads = Zygote.gradient((m,x,y) -> (m(x) - y)^2, model, data[1]...)
model.w .-= 0.1 .* grads[1].v # grads[1] has same fields as model
```
Instead of gradient descent, we can use say Adam, applied to all the parameters, like this – Optimisers.jl again understands structs with matching fields:
```julia
optstate = Optimisers.setup(Adam(), model)
Optimisers.update!(optstate, model, grads[1]);
```
That covers 2 of your 3 bullet points, and we haven’t loaded Flux.jl yet. Many loss functions are so simple you can just write them out, like `(m(x) - y)^2`. But Flux.jl is one place to get a library of standard ones. (Plus standard model-building layers, which it sounds like you may not want.)
Instead of the “anonymous struct” created by `x -> ...`, you can make your own `struct MyModel; w::Vector{Float32}; ...`, and make it callable `(m::MyModel)(x) = sum(m.v .* tanh.(m.w*x .+ m.b))`. Fluxperimental.jl has some tools for making this a little easer (and easier to revise) but using basic Julia is fine too.
---
### Author:  [@stur86](https://discourse.julialang.org/u/stur86)
#### Post date: [April 21, 2025, 8:27pm UTC](https://discourse.julialang.org/t/best-libraries-for-experimenting-with-small-scale-models/128106/4 "2025-04-21T20:27:10Z")
Yeah, I was considering something like this. Simple loss functions like square error or cross-entropy are no problem, but it’s nice to have standard implementations regardless whenever possible.
---
### Author:  [@mcabbott](https://discourse.julialang.org/u/mcabbott)
#### Post date: [April 21, 2025, 9:00pm UTC](https://discourse.julialang.org/t/best-libraries-for-experimenting-with-small-scale-models/128106/5 "2025-04-21T21:00:13Z")
It would not be crazy to split [these functions](https://github.com/FluxML/Flux.jl/blob/master/src/losses/functions.jl) out into their own little package, if someone wants it enough. Several such things have been removed, e.g. Tracker.jl, OneHotArrays.jl, arguably Optimisers.jl
---
### Author:  [@digital\_carver](https://discourse.julialang.org/u/digital_carver)
#### Post date: [April 21, 2025, 10:35pm UTC](https://discourse.julialang.org/t/best-libraries-for-experimenting-with-small-scale-models/128106/6 "2025-04-21T22:35:37Z")
I have no idea how (or whether) it fits the rest of your requirements, but “small-ish models” and “running on CPU” brought to mind SimpleChains.jl: [Doing small network scientific machine learning in Julia 5x faster than PyTorch](https://julialang.org/blog/2022/04/simple-chains/)
---
### Author:  [@DoktorMike](https://discourse.julialang.org/u/DoktorMike)
#### Post date: [April 23, 2025, 6:42pm UTC](https://discourse.julialang.org/t/best-libraries-for-experimenting-with-small-scale-models/128106/7 "2025-04-23T18:42:31Z")
I think that would in general be a good idea. I don’t really see a reason why all the deep learning frameworks shouldn’t share the implementation of the loss functions. But I might be missing something.
---
### Author:  [@jokctseng](https://discourse.julialang.org/u/jokctseng)
#### Post date: [April 25, 2025, 2:38am UTC](https://discourse.julialang.org/t/best-libraries-for-experimenting-with-small-scale-models/128106/8 "2025-04-25T02:38:27Z")
This post was temporarily hidden by the community for possibly being off-topic, unfocused, inappropriate, or spammy.