Models

FluxTraining.jl works with all Flux . jl -compatible models. Unless you are using a custom training loop , a model is expected to take a single input xs , which corresponds to the encoded inputs returned by your data iterator . This means the following has to work:

			
			
			
			
			xs
			
			,
			
			 
			
			ys
			
			 
			
			=
			
			 
			
			
			first
			
			(
			
			dataiter
			
			)
			
			

			
			
			ŷs
			
			 
			
			=
			
			 
			
			
			model
			
			(
			
			xs
			
			)

model also has to be differentiable. If you’re composing Flux.jl layers, this is likely the case. You can always make sure by testing:

			
			
			
			using
			
			
			 
			
			Flux
			
			,
			
			
			 
			
			Zygote
			
			

			
			

			
			
			
			xs
			
			,
			
			 
			
			ys
			
			 
			
			=
			
			 
			
			
			first
			
			(
			
			dataiter
			
			)
			
			

			
			
			lossfn
			
			 
			
			=
			
			 
			
			
			Flux
			
			.
			
			
			mse
			
			

			
			
			grads
			
			 
			
			=
			
			 
			
			
			
			
			Zygote
			
			.
			
			
			gradient
			
			(
			
			
			
			Flux
			
			.
			
			
			params
			
			(
			
			model
			
			)
			
			)
			
			 
			
			do
			
			
			
			
    
			
			
			
			lossfn
			
			(
			
			
			model
			
			(
			
			xs
			
			)
			
			,
			
			 
			
			ys
			
			)
			
			

			
			end

Creating models

The simplest way to create a Flux.jl-compatible model is to use layers from Flux.jl.
A good entrypoint is this tutorial in Flux’s documentation.

There is also a large number of packages that provide complete model architectures or domain-specific layers. Below is a non-exhaustive list:

• Metalhead . jl implements common model architectures for computer vision,

• GraphNeuralNetworks . jl provides layers and utilities for graph neural networks,

• Transformers . jl implements transformer models including pretrained language models