repository·dev·Indexed 41 days ago
https://github.com/fluxml/mljflux.jlAn interface that integrates Flux deep learning models into the MLJ machine learning framework, enabling standardized workflows for deep learning in Julia. The library provides support for architecture search, model composition, entity embeddings, hyperparameter tuning, incremental training, and live training.
MLJFlux.jl is an interface that connects the MLJ machine learning toolbox with the Flux.jl deep learning framework in Julia. It allows users to implement supervised deep learning models using the Flux syntax while adhering to the standard MLJ workflow.
IteratedModels, you gain access to advanced training controls such as early stopping, model snapshots, callbacks, and cyclic learning rates.MLJ machines cache state, allowing for "warm restarts" of model training. When calling fit!(mach), a warm restart is automatically used if:
model.epochs has changed since the last call.model.epochs or model.optimiser have changed, AND model.optimiser_changes_trigger_retraining == false (the default). This allows you to dynamically modify learning rates without losing learned weights.Note that the "state" part of the optimiser is ignored in the comparison for the second condition. You can access the associated MLJ model via model = mach.model.
Iterative MLJFlux models can be wrapped using MLJ.IteratedModel to implement sophisticated stopping criteria and monitoring. This is useful for preventing overfitting or stopping when training plateaus.
Common Controls:
Step(n): Apply checks every n epochs.Patience(n): Stop after n consecutive increases in loss.InvalidValue(): Stop if loss becomes NaN or Inf.TimeLimit(t): Stop after t time has elapsed.Save(path, predicate=n): Save a snapshot of the machine every n control cycles.WithLossDo(callback): Execute a callback with the current loss.WithTrainingLossesDo(callback): Execute a callback with the current training loss.WithIterationsDo(callback): Execute a callback with the current epoch number.Workflow:
controls.iterated_clf = IteratedModel(clf, controls=controls, ...).fit!.using MLJ
using MLJIteration
# Example controls setup
controls = [
Step(2),
Patience(3),
InvalidValue(),
TimeLimit(5/60),
MLJIteration.skip(Save(joinpath(tempdir(), "mnist.jls")), predicate=3),
WithLossDo(update_loss),
WithTrainingLossesDo(update_training_loss),
WithIterationsDo(update_epochs)
]
# Wrap the classifier
iterated_clf = IteratedModel(
clf,
controls=controls,
resampling=Holdout(fraction_train=0.7),
measure=log_loss,
)
# Train the wrapped model
mach = machine(iterated_clf, images, labels)
fit!(mach, rows=train)MLJ follows the convention that rows = observations. However, Flux (the underlying engine) uses the opposite convention (columns = observations).
If your data is a matrix where the column index represents the observation, you must present the adjoint or transpose of your matrix to MLJFlux models to ensure correct training. You can also use permutedims to transform the matrix before passing it to the model.
To use a Flux model within the MLJ ecosystem, you must wrap the neural network architecture in one of the four provided MLJFlux containers.
Each container uses a builder attribute to define the neural network architecture (the Flux model) based on the shape of the input data. The containers are categorized by the type of task:
MLJModelInterface.Deterministic.MLJModelInterface.Probabilistic.By wrapping your Flux model, you can specify hyperparameters such as the loss function, optimizer, and number of epochs, making the deep learning model behave like any other standard MLJ model.
Most MLJFlux models support categorical features by learning entity embeddings. To use these learned embeddings in standard MLJ pipelines—specifically for supervised models that are not neural networks and require features to be of type Continuous—wrap your MLJFlux model with MLJFlux.EntityEmbedder. This allows the embeddings to function as transformer elements within an MLJ pipeline.
# Wrap an MLJFlux model to use its learned embeddings as continuous features in an MLJ pipeline
wrapped_model = MLJFlux.EntityEmbedder(mljflux_model)The incremental_training directory contains various formats of the incremental training workflow for testing and demonstration purposes:
notebook.ipynb (executed) and notebook.unexecuted.ipynb (unexecuted) for interactive exploration.notebook.md which is integrated into the official MLJFlux.jl documentation.notebook.jl which is an executable Julia script containing annotated comments explaining the workflow.Note: To execute these scripts or notebooks, you must have the project environment configured. The required Manifest.toml and Project.toml files are located at ../../../docs/ relative to this directory.
# To run the executable script version:
include("notebook.jl")To tune hyperparameters in MLJFlux, you can use the TunedModel constructor from MLJ. This involves defining hyperparameter ranges using the range function, specifying a tuning strategy (like Grid), and a resampling method (like CV for cross-validation).
Note that MLJFlux models often require data to be in Float32 format for compatibility with Flux.jl. When defining ranges for optimisers, you should pass instances of optimisers (e.g., from the Optimisers.jl package) rather than just learning rate scalars.
# 1. Instantiate the model
NeuralNetworkClassifier = @load NeuralNetworkClassifier pkg=MLJFlux
clf = NeuralNetworkClassifier(
builder=MLJFlux.MLP(; hidden=(5,4), σ=Flux.relu),
optimiser=Optimisers.Adam(0.01),
batch_size=8,
epochs=10,
rng=stable_rng(),
)
# 2. Define hyperparameter ranges
r1 = range(clf, :batch_size, lower=1, upper=64)
etas = [10^x for x in range(-4, stop=0, length=4)]
optimisers = [Optimisers.Adam(eta) for eta in etas]
r2 = range(clf, :optimiser, values=optimisers)
# 3. Create the TunedModel
tuned_model = TunedModel(
model=clf,
tuning=Grid(goal=25),
resampling=CV(nfolds=5, rng=stable_rng()),
range=[r1, r2],
measure=cross_entropy,
)
# 4. Fit the machine
mach = machine(tuned_model, X, y)
fit!(mach, verbosity=0)
# 5. Access the best model
best_model = fitted_params(mach).best_modelWhen working with MLJFlux for incremental training, ensure you have the following packages imported. Note that Optimisers.jl should be used instead of the native Flux.jl optimisers, as the latter are no longer supported.
Required imports:
MLJ: To access MLJFlux models.Flux: For model flexibility (e.g., activation functions).Optimisers: For model optimisers.StableRNGs: To ensure reproducibility across Julia versions.using MLJ # Has MLJFlux models
using Flux # For more flexibility
import Optimisers # native Flux.jl optimisers no longer supported
using StableRNGs # for reproducibility across Julia versions
stable_rng() = StableRNGs.StableRNG(123)The entity embeddings workflow is demonstrated through several file formats in the docs/src/common_workflows/entity_embeddings/ directory. You can use the following files to understand and implement the workflow:
notebook.ipynb: An executed Jupyter notebook showing the full workflow results.notebook.unexecuted.ipynb: An unexecuted Jupyter notebook for local experimentation.notebook.md: A static markdown version included in the official MLJFlux.jl documentation.notebook.jl: An executable Julia script containing annotated code comments for step-by-step execution.To ensure the MNIST examples execute reliably, you must include the accompanying Manifest.toml and Project.toml files in your environment. This ensures all dependency versions match the environment for which the examples were written.
Available files in this directory:
notebook.ipynb: An executed Jupyter notebook.notebook.unexecuted.ipynb: An unexecuted Jupyter notebook.notebook.md: Static markdown version of the notebook.notebook.jl: An executable Julia script with annotated comments.