Merge pull request #25 from JuliaTrustworthyAI/lts

done

.github/workflows/CI.yml

@@ -18,9 +18,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.7'
-          - '1.8'
-          - '1.9'
+          - '1.6'
           - '1.10'
         os:
           - ubuntu-latest

CHANGELOG.md

@@ -6,6 +6,12 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 *Note*: We try to adhere to these practices as of version [v0.1.4].
 
+## Version [0.1.5] - 2024-06-07
+
+### Changed
+
+- Added support for LTS (1.6). [#25]
+
 ## Version [0.1.4] - 2024-06-04
 
 ### Changed

Project.toml

@@ -1,7 +1,7 @@
 name = "JointEnergyModels"
 uuid = "48c56d24-211d-4463-bbc0-7a701b291131"
 authors = ["Patrick Altmeyer"]
-version = "0.1.4"
+version = "0.1.5"
 
 [deps]
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"
@@ -31,14 +31,14 @@ MLJFlux = "0.2, 0.3, 0.4.0"
 MLJModelInterface = "1.8"
 MLUtils = "0.4"
 ProgressMeter = "1.7"
+Random = "1.6, 1.10"
 Reexport = "1.2.2"
 StatsBase = "0.33, 0.34"
 Tables = "1.10"
 TaijaBase = "1.1.0"
+Test = "1.6, 1.10"
 Zygote = "0.6"
-Random = "1.7, 1.10"
-Test = "1.7, 1.10"
-julia = "1.7, 1.10"
+julia = "1.6, 1.10"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"

docs/assets.jl

@@ -1,2 +1,2 @@
 # Since Documenter.jl cannot handle links to anything outside the docs/src directory, we need to copy the LICENSE and CHANGELOG.md files to the src/assets directory: https://discourse.julialang.org/t/make-documenter-jl-understand-where-to-search-local-markdown-links/84012/6?u=pat-alt
-cp(joinpath(@__DIR__, "..", "LICENSE"), joinpath(@__DIR__, "src/LICENSE"); force=true)
+cp(joinpath(@__DIR__, "..", "LICENSE"), joinpath(@__DIR__, "src/LICENSE"); force = true)

docs/make.jl

@@ -5,20 +5,20 @@ include("setup_docs.jl")
 
 include("assets.jl")
 
-DocMeta.setdocmeta!(JointEnergyModels, :DocTestSetup, :(setup_docs); recursive=true)
+DocMeta.setdocmeta!(JointEnergyModels, :DocTestSetup, :(setup_docs); recursive = true)
 
 makedocs(;
-    modules=[JointEnergyModels],
-    authors="Patrick Altmeyer",
-    repo="https://github.com/JuliaTrustworthyAI/JointEnergyModels.jl/blob/{commit}{path}#{line}",
-    sitename="JointEnergyModels.jl",
-    format=Documenter.HTML(;
-        prettyurls=get(ENV, "CI", "false") == "true",
-        canonical="https://juliatrustworthyai.github.io/JointEnergyModels.jl",
-        edit_link="main",
-        assets=String[]
+    modules = [JointEnergyModels],
+    authors = "Patrick Altmeyer",
+    repo = "https://github.com/JuliaTrustworthyAI/JointEnergyModels.jl/blob/{commit}{path}#{line}",
+    sitename = "JointEnergyModels.jl",
+    format = Documenter.HTML(;
+        prettyurls = get(ENV, "CI", "false") == "true",
+        canonical = "https://juliatrustworthyai.github.io/JointEnergyModels.jl",
+        edit_link = "main",
+        assets = String[],
     ),
-    pages=[
+    pages = [
         "🏠 Home" => "index.md",
         # "🫣 Tutorials" => [
         #     "Overview" => "tutorials/index.md",
@@ -32,10 +32,10 @@ makedocs(;
         # "🧐 Reference" => "reference.md",
         # "🛠 Contribute" => "contribute.md",
         # "📚 Additional Resources" => "assets/resources.md",
-    ]
+    ],
 )
 
 deploydocs(;
-    repo="github.com/JuliaTrustworthyAI/JointEnergyModels.jl",
-    devbranch="main",
+    repo = "github.com/JuliaTrustworthyAI/JointEnergyModels.jl",
+    devbranch = "main",
 )

docs/src/utils/mnist.jl

@@ -1,4 +1,4 @@
-function _resize(x; size=(28, 28))
+function _resize(x; size = (28, 28))
     if n_digits != 28
         img_source = MLDatasets.convert2image(MNIST, x)
         img_rs = imresize(img_source, size)
@@ -7,25 +7,25 @@ function _resize(x; size=(28, 28))
     return x
 end
 
-function pre_process(x; noise::Float32=0.03f0)
+function pre_process(x; noise::Float32 = 0.03f0)
     ϵ = Float32.(randn(size(x)) * noise)
     x = @.(2 * x - 1) .+ ϵ
     return x
 end
 
-function load_mnist_data(; nobs::Int=1000, n_digits::Int=1000)
+function load_mnist_data(; nobs::Int = 1000, n_digits::Int = 1000)
     # Train Set:
-    Xtrain, ytrain = MNIST(split=:train)[:]
+    Xtrain, ytrain = MNIST(split = :train)[:]
     end_train = minimum([nobs, size(Xtrain)[end]])
     Xtrain = Xtrain[:, :, 1:end_train]
-    Xtrain = mapslices(x -> _resize(x; size=(n_digits, n_digits)), Xtrain, dims=(1, 2))
+    Xtrain = mapslices(x -> _resize(x; size = (n_digits, n_digits)), Xtrain, dims = (1, 2))
     ytrain = ytrain[1:end_train]
 
     # Test Set:
-    Xtest, ytest = MNIST(split=:test)[:]
+    Xtest, ytest = MNIST(split = :test)[:]
     end_test = minimum([nobs, size(Xtest)[end]])
     Xtest = Xtest[:, :, 1:end_test]
-    Xtest = mapslices(x -> _resize(x; size=(n_digits, n_digits)), Xtest, dims=(1, 2))
+    Xtest = mapslices(x -> _resize(x; size = (n_digits, n_digits)), Xtest, dims = (1, 2))
     ytest = ytest[1:end_test]
 
     ## One-hot-encode the labels
@@ -34,26 +34,49 @@ function load_mnist_data(; nobs::Int=1000, n_digits::Int=1000)
     ## Validation Set:
     num_val = Int(round(nobs / 10))
     Xtrain, Xval = (Xtrain[:, :, 1:(end-num_val)], Xtrain[:, :, (end-num_val+1):end])
-    Xtrain = mapslices(x -> pre_process(x), Xtrain, dims=(1, 2))
-    Xval = mapslices(x -> pre_process(x, noise=0.0f0), Xval, dims=(1, 2))
+    Xtrain = mapslices(x -> pre_process(x), Xtrain, dims = (1, 2))
+    Xval = mapslices(x -> pre_process(x, noise = 0.0f0), Xval, dims = (1, 2))
     ytrain, yval = (ytrain[:, 1:(end-num_val)], ytrain[:, (end-num_val+1):end])
 
     return Xtrain, ytrain, Xval, yval, Xtest, ytest
 end
 
-function samples_real(model::JointEnergyModel, dl::DataLoader, n::Int=16; img_size=n_digits * 5)
-    x = reduce((x, y) -> cat(x, y[1], dims=ndims(x)), dl, init=[])
+function samples_real(
+    model::JointEnergyModel,
+    dl::DataLoader,
+    n::Int = 16;
+    img_size = n_digits * 5,
+)
+    x = reduce((x, y) -> cat(x, y[1], dims = ndims(x)), dl, init = [])
     num_x = Int(round(ceil(sqrt(n))))
     num_y = Int(round(floor(sqrt(n))))
-    plot_data = [heatmap(rotl90(x[:, :, rand(1:size(x)[end])]), axis=nothing, cb=false) for i in 1:n]
-    plot(plot_data..., layout=(num_x, num_y), size=(num_x * img_size, num_y * img_size), margin=(round(0.05 * img_size), :px))
+    plot_data = [
+        heatmap(rotl90(x[:, :, rand(1:size(x)[end])]), axis = nothing, cb = false) for
+        i = 1:n
+    ]
+    plot(
+        plot_data...,
+        layout = (num_x, num_y),
+        size = (num_x * img_size, num_y * img_size),
+        margin = (round(0.05 * img_size), :px),
+    )
 end
 
-function samples_generated(model::JointEnergyModel, dl::DataLoader, n::Int=16; img_size=n_digits * 5)
-    x = reduce((x, y) -> cat(x, y[1], dims=ndims(x)), dl, init=[])
+function samples_generated(
+    model::JointEnergyModel,
+    dl::DataLoader,
+    n::Int = 16;
+    img_size = n_digits * 5,
+)
+    x = reduce((x, y) -> cat(x, y[1], dims = ndims(x)), dl, init = [])
     x = jem.sampler(jem.chain, jem.sampling_rule, size(x))
     num_x = Int(round(ceil(sqrt(n))))
     num_y = Int(round(floor(sqrt(n))))
-    plot_data = [heatmap(rotl90(x[:, :, i]), axis=nothing, cb=false) for i in 1:n]
-    plot(plot_data..., layout=(num_x, num_y), size=(num_x * img_size, num_y * img_size), margin=(round(0.05 * img_size), :px))
+    plot_data = [heatmap(rotl90(x[:, :, i]), axis = nothing, cb = false) for i = 1:n]
+    plot(
+        plot_data...,
+        layout = (num_x, num_y),
+        size = (num_x * img_size, num_y * img_size),
+        margin = (round(0.05 * img_size), :px),
+    )
 end

docs/src/utils/utils.jl

@@ -1,2 +1,2 @@
 # MNIST
-include("mnist.jl")
+include("mnist.jl")

src/mlj_flux.jl

@@ -6,7 +6,7 @@ using ProgressMeter
 using Random
 using Tables
 
-const default_builder_jem = MLJFlux.MLP(hidden=(32, 32, 32,), σ=Flux.swish)
+const default_builder_jem = MLJFlux.MLP(hidden = (32, 32, 32), σ = Flux.swish)
 
 "The `JointEnergyClassifier` struct is a wrapper for a `JointEnergyModel` that can be used with MLJFlux.jl."
 mutable struct JointEnergyClassifier{B,F,O,L} <: MLJFlux.MLJFluxProbabilistic
@@ -29,22 +29,38 @@ end
 
 function JointEnergyClassifier(
     sampler::AbstractSampler;
-    builder::B=default_builder_jem,
-    finaliser::F=Flux.softmax,
-    optimiser::O=Flux.Optimise.Adam(),
-    loss::L=Flux.crossentropy,
-    epochs::Int=100, batch_size::Int=100, lambda::Float64=0.0, alpha::Float64=0.0,
-    rng::Union{AbstractRNG,Int64}=Random.GLOBAL_RNG,
-    optimiser_changes_trigger_retraining::Bool=false,
-    acceleration::AbstractResource=CPU1(),
-    jem_training_params::NamedTuple=(verbosity=epochs,num_epochs=epochs,),
-    sampling_steps::Union{Nothing,Int}=nothing,
+    builder::B = default_builder_jem,
+    finaliser::F = Flux.softmax,
+    optimiser::O = Flux.Optimise.Adam(),
+    loss::L = Flux.crossentropy,
+    epochs::Int = 100,
+    batch_size::Int = 100,
+    lambda::Float64 = 0.0,
+    alpha::Float64 = 0.0,
+    rng::Union{AbstractRNG,Int64} = Random.GLOBAL_RNG,
+    optimiser_changes_trigger_retraining::Bool = false,
+    acceleration::AbstractResource = CPU1(),
+    jem_training_params::NamedTuple = (verbosity = epochs, num_epochs = epochs),
+    sampling_steps::Union{Nothing,Int} = nothing,
 ) where {B,F,O,L}
 
     # Initialise the MLJFlux wrapper:
     mlj_jem = JointEnergyClassifier(
-        builder, finaliser, optimiser, loss, epochs, batch_size, lambda, alpha, rng,
-        optimiser_changes_trigger_retraining, acceleration, sampler, nothing, jem_training_params, sampling_steps
+        builder,
+        finaliser,
+        optimiser,
+        loss,
+        epochs,
+        batch_size,
+        lambda,
+        alpha,
+        rng,
+        optimiser_changes_trigger_retraining,
+        acceleration,
+        sampler,
+        nothing,
+        jem_training_params,
+        sampling_steps,
     )
 
     return mlj_jem
@@ -71,16 +87,10 @@ function MLJFlux.build(model::JointEnergyClassifier, rng, shape)
 
     # JointEnergyModel:
     if isnothing(model.sampling_steps)
-        model.jem = JointEnergyModel(
-            chain,
-            model.sampler,
-        )
+        model.jem = JointEnergyModel(chain, model.sampler)
     else
-        model.jem = JointEnergyModel(
-            chain,
-            model.sampler;
-            sampling_steps=model.sampling_steps,
-        )
+        model.jem =
+            JointEnergyModel(chain, model.sampler; sampling_steps = model.sampling_steps)
     end
 
     return chain
@@ -89,46 +99,60 @@ end
 # returns the model `fitresult` (see "Adding Models for General Use"
 # section of the MLJ manual) which must always have the form `(chain,
 # metadata)`, where `metadata` is anything extra needed by `predict`:
-MLJFlux.fitresult(model::JointEnergyClassifier, chain, y) =
-    (chain, MMI.classes(y[1]))
+MLJFlux.fitresult(model::JointEnergyClassifier, chain, y) = (chain, MMI.classes(y[1]))
 
-function MMI.predict(
-    model::JointEnergyClassifier,
-    fitresult,
-    Xnew
-)
+function MMI.predict(model::JointEnergyClassifier, fitresult, Xnew)
     chain, levels = fitresult
     X = MLJFlux.reformat(Xnew)
-    probs = vcat([chain(MLJFlux.tomat(X[:, i]))' for i in 1:size(X, 2)]...)
+    probs = vcat([chain(MLJFlux.tomat(X[:, i]))' for i = 1:size(X, 2)]...)
     return MMI.UnivariateFinite(levels, probs)
 end
 
-MMI.metadata_model(JointEnergyClassifier,
-    input=Union{AbstractArray,MMI.Table(MMI.Continuous)},
-    target=AbstractVector{<:MMI.Finite},
-    path="MLJFlux.JointEnergyClassifier")
+MMI.metadata_model(
+    JointEnergyClassifier,
+    input = Union{AbstractArray,MMI.Table(MMI.Continuous)},
+    target = AbstractVector{<:MMI.Finite},
+    path = "MLJFlux.JointEnergyClassifier",
+)
 
-function MLJFlux.fit!(model::JointEnergyClassifier, penalty, chain, optimiser, epochs, verbosity, X, y)
+function MLJFlux.fit!(
+    model::JointEnergyClassifier,
+    penalty,
+    chain,
+    optimiser,
+    epochs,
+    verbosity,
+    X,
+    y,
+)
 
     loss = model.loss
 
     # intitialize and start progress meter:
-    meter = Progress(epochs, dt=0, desc="Optimising neural net:",
-        barglyphs=BarGlyphs("[=> ]"), barlen=25, color=:yellow)
+    meter = Progress(
+        epochs,
+        dt = 0,
+        desc = "Optimising neural net:",
+        barglyphs = BarGlyphs("[=> ]"),
+        barlen = 25,
+        color = :yellow,
+    )
     verbosity != 1 || next!(meter)
 
     # initiate training:
     train_set = zip(X, y)
     opt_state = Flux.setup(optimiser, model.jem)
 
     history = train_model(
-        model.jem, train_set, opt_state;
-        class_loss_fun=loss,
-        progress_meter=meter,
-        num_epochs=model.epochs,
+        model.jem,
+        train_set,
+        opt_state;
+        class_loss_fun = loss,
+        progress_meter = meter,
+        num_epochs = model.epochs,
         model.jem_training_params...,
     )
 
     return model.jem.chain, history
 
-end
+end