Convolutional VAE for MNIST
Convolutional variational autoencoder (CVAE) implementation in MLX using MNIST. This is based on the CVAE implementation in MLX.
julia
using Lux,
Reactant,
MLDatasets,
Random,
Statistics,
Enzyme,
MLUtils,
DataAugmentation,
ConcreteStructs,
OneHotArrays,
ImageShow,
Images,
Printf,
Optimisers
const xdev = reactant_device(; force=true)
const cdev = cpu_device()
const IN_VSCODE = isdefined(Main, :VSCodeServer)falseModel Definition
First we will define the encoder.It maps the input to a normal distribution in latent space and sample a latent vector from that distribution.
julia
function cvae_encoder(
rng=Random.default_rng();
num_latent_dims::Int,
image_shape::Dims{3},
max_num_filters::Int,
)
flattened_dim = prod(image_shape[1:2] .÷ 8) * max_num_filters
return @compact(;
embed=Chain(
Chain(
Conv((3, 3), image_shape[3] => max_num_filters ÷ 4; stride=2, pad=1),
BatchNorm(max_num_filters ÷ 4, leakyrelu),
),
Chain(
Conv((3, 3), max_num_filters ÷ 4 => max_num_filters ÷ 2; stride=2, pad=1),
BatchNorm(max_num_filters ÷ 2, leakyrelu),
),
Chain(
Conv((3, 3), max_num_filters ÷ 2 => max_num_filters; stride=2, pad=1),
BatchNorm(max_num_filters, leakyrelu),
),
FlattenLayer(),
),
proj_mu=Dense(flattened_dim, num_latent_dims; init_bias=zeros32),
proj_log_var=Dense(flattened_dim, num_latent_dims; init_bias=zeros32),
rng
) do x
y = embed(x)
μ = proj_mu(y)
logσ² = proj_log_var(y)
T = eltype(logσ²)
logσ² = clamp.(logσ², -T(20.0f0), T(10.0f0))
σ = exp.(logσ² .* T(0.5))
# Generate a tensor of random values from a normal distribution
ϵ = randn_like(Lux.replicate(rng), σ)
# Reparameterization trick to backpropagate through sampling
z = ϵ .* σ .+ μ
@return z, μ, logσ²
end
endSimilarly we define the decoder.
julia
function cvae_decoder(; num_latent_dims::Int, image_shape::Dims{3}, max_num_filters::Int)
flattened_dim = prod(image_shape[1:2] .÷ 8) * max_num_filters
return @compact(;
linear=Dense(num_latent_dims, flattened_dim),
upchain=Chain(
Chain(
Upsample(2),
Conv((3, 3), max_num_filters => max_num_filters ÷ 2; stride=1, pad=1),
BatchNorm(max_num_filters ÷ 2, leakyrelu),
),
Chain(
Upsample(2),
Conv((3, 3), max_num_filters ÷ 2 => max_num_filters ÷ 4; stride=1, pad=1),
BatchNorm(max_num_filters ÷ 4, leakyrelu),
),
Chain(
Upsample(2),
Conv(
(3, 3), max_num_filters ÷ 4 => image_shape[3], sigmoid; stride=1, pad=1
),
),
),
max_num_filters
) do x
y = linear(x)
img = reshape(y, image_shape[1] ÷ 8, image_shape[2] ÷ 8, max_num_filters, :)
@return upchain(img)
end
end
@concrete struct CVAE <: AbstractLuxContainerLayer{(:encoder, :decoder)}
encoder <: AbstractLuxLayer
decoder <: AbstractLuxLayer
end
function CVAE(
rng=Random.default_rng();
num_latent_dims::Int,
image_shape::Dims{3},
max_num_filters::Int,
)
decoder = cvae_decoder(; num_latent_dims, image_shape, max_num_filters)
encoder = cvae_encoder(rng; num_latent_dims, image_shape, max_num_filters)
return CVAE(encoder, decoder)
end
function (cvae::CVAE)(x, ps, st)
(z, μ, logσ²), st_enc = cvae.encoder(x, ps.encoder, st.encoder)
x_rec, st_dec = cvae.decoder(z, ps.decoder, st.decoder)
return (x_rec, μ, logσ²), (; encoder=st_enc, decoder=st_dec)
end
function encode(cvae::CVAE, x, ps, st)
(z, _, _), st_enc = cvae.encoder(x, ps.encoder, st.encoder)
return z, (; encoder=st_enc, st.decoder)
end
function decode(cvae::CVAE, z, ps, st)
x_rec, st_dec = cvae.decoder(z, ps.decoder, st.decoder)
return x_rec, (; decoder=st_dec, st.encoder)
endLoading MNIST
julia
@concrete struct TensorDataset
dataset
transform
total_samples::Int
end
Base.length(ds::TensorDataset) = ds.total_samples
function Base.getindex(ds::TensorDataset, idxs::Union{Vector{<:Integer},AbstractRange})
img = Image.(eachslice(convert2image(ds.dataset, idxs); dims=3))
return stack(parent ∘ itemdata ∘ Base.Fix1(apply, ds.transform), img)
end
function loadmnist(batchsize, image_size::Dims{2})
# Load MNIST: Only 1500 for demonstration purposes on CI
train_dataset = MNIST(; split=:train)
N = parse(Bool, get(ENV, "CI", "false")) ? 5000 : length(train_dataset)
train_transform = ScaleKeepAspect(image_size) |> ImageToTensor()
trainset = TensorDataset(train_dataset, train_transform, N)
trainloader = DataLoader(trainset; batchsize, shuffle=true, partial=false)
return trainloader
endHelper Functions
Generate an Image Grid from a list of images
julia
function create_image_grid(imgs::AbstractArray, grid_rows::Int, grid_cols::Int)
total_images = grid_rows * grid_cols
imgs = map(eachslice(imgs[:, :, :, 1:total_images]; dims=4)) do img
cimg = if size(img, 3) == 1
colorview(Gray, view(img, :, :, 1))
else
colorview(RGB, permutedims(img, (3, 1, 2)))
end
return cimg'
end
return create_image_grid(imgs, grid_rows, grid_cols)
end
function create_image_grid(images::Vector, grid_rows::Int, grid_cols::Int)
# Check if the number of images matches the grid
total_images = grid_rows * grid_cols
@assert length(images) == total_images
# Get the size of a single image (assuming all images are the same size)
img_height, img_width = size(images[1])
# Create a blank grid canvas
grid_height = img_height * grid_rows
grid_width = img_width * grid_cols
grid_canvas = similar(images[1], grid_height, grid_width)
# Place each image in the correct position on the canvas
for idx in 1:total_images
row = div(idx - 1, grid_cols) + 1
col = mod(idx - 1, grid_cols) + 1
start_row = (row - 1) * img_height + 1
start_col = (col - 1) * img_width + 1
grid_canvas[start_row:(start_row + img_height - 1), start_col:(start_col + img_width - 1)] .= images[idx]
end
return grid_canvas
end
function loss_function(model, ps, st, X)
(y, μ, logσ²), st = model(X, ps, st)
reconstruction_loss = MSELoss(; agg=sum)(y, X)
kldiv_loss = -sum(1 .+ logσ² .- μ .^ 2 .- exp.(logσ²)) / 2
loss = reconstruction_loss + kldiv_loss
return loss, st, (; y, μ, logσ², reconstruction_loss, kldiv_loss)
end
function generate_images(
model, ps, st; num_samples::Int=128, num_latent_dims::Int, decode_compiled=nothing
)
z = get_device((ps, st))(randn(Float32, num_latent_dims, num_samples))
if decode_compiled === nothing
images, _ = decode(model, z, ps, Lux.testmode(st))
else
images, _ = decode_compiled(model, z, ps, Lux.testmode(st))
images = cpu_device()(images)
end
return create_image_grid(images, 8, num_samples ÷ 8)
end
function reconstruct_images(model, ps, st, X)
(recon, _, _), _ = model(X, ps, Lux.testmode(st))
recon = cpu_device()(recon)
return create_image_grid(recon, 8, size(X, ndims(X)) ÷ 8)
endreconstruct_images (generic function with 1 method)Training the Model
julia
function main(;
batchsize=128,
image_size=(64, 64),
num_latent_dims=8,
max_num_filters=64,
seed=0,
epochs=50,
weight_decay=1.0e-5,
learning_rate=1.0e-3,
num_samples=batchsize,
)
rng = Xoshiro()
Random.seed!(rng, seed)
cvae = CVAE(rng; num_latent_dims, image_shape=(image_size..., 1), max_num_filters)
ps, st = xdev(Lux.setup(rng, cvae))
z = xdev(randn(Float32, num_latent_dims, num_samples))
decode_compiled = Reactant.with_config(;
dot_general_precision=PrecisionConfig.HIGH,
convolution_precision=PrecisionConfig.HIGH,
) do
@compile decode(cvae, z, ps, Lux.testmode(st))
end
x = xdev(randn(Float32, image_size..., 1, batchsize))
cvae_compiled = Reactant.with_config(;
dot_general_precision=PrecisionConfig.HIGH,
convolution_precision=PrecisionConfig.HIGH,
) do
@compile cvae(x, ps, Lux.testmode(st))
end
train_dataloader = xdev(loadmnist(batchsize, image_size))
opt = AdamW(; eta=learning_rate, lambda=weight_decay)
train_state = Training.TrainState(cvae, ps, st, opt)
@printf "Total Trainable Parameters: %0.4f M\n" (Lux.parameterlength(ps) / 1.0e6)
empty_row, model_img_full = nothing, nothing
for epoch in 1:epochs
loss_total = 0.0f0
total_samples = 0
start_time = time()
for (i, X) in enumerate(train_dataloader)
(_, loss, _, train_state) = Training.single_train_step!(
AutoEnzyme(), loss_function, X, train_state; return_gradients=Val(false)
)
loss_total += loss
total_samples += size(X, ndims(X))
if i % 250 == 0 || i == length(train_dataloader)
throughput = total_samples / (time() - start_time)
@printf "Epoch %d, Iter %d, Loss: %.7f, Throughput: %.6f im/s\n" epoch i loss throughput
end
end
total_time = time() - start_time
train_loss = loss_total / length(train_dataloader)
throughput = total_samples / total_time
@printf "Epoch %d, Train Loss: %.7f, Time: %.4fs, Throughput: %.6f im/s\n" epoch train_loss total_time throughput
if IN_VSCODE || epoch == epochs
recon_images = reconstruct_images(
cvae_compiled,
train_state.parameters,
train_state.states,
first(train_dataloader),
)
gen_images = generate_images(
cvae,
train_state.parameters,
train_state.states;
num_samples,
num_latent_dims,
decode_compiled,
)
if empty_row === nothing
empty_row = similar(gen_images, image_size[1], size(gen_images, 2))
fill!(empty_row, 0)
end
model_img_full = vcat(recon_images, empty_row, gen_images)
IN_VSCODE && display(model_img_full)
end
end
return model_img_full
end
img = main()WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1760328583.441699 1460684 service.cc:158] XLA service 0x1c95c060 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
I0000 00:00:1760328583.441757 1460684 service.cc:166] StreamExecutor device (0): Quadro RTX 5000, Compute Capability 7.5
I0000 00:00:1760328583.442744 1460684 se_gpu_pjrt_client.cc:1339] Using BFC allocator.
I0000 00:00:1760328583.442786 1460684 gpu_helpers.cc:136] XLA backend allocating 12526534656 bytes on device 0 for BFCAllocator.
I0000 00:00:1760328583.442846 1460684 gpu_helpers.cc:177] XLA backend will use up to 4175511552 bytes on device 0 for CollectiveBFCAllocator.
I0000 00:00:1760328583.452061 1460684 cuda_dnn.cc:463] Loaded cuDNN version 91200
Total Trainable Parameters: 0.1493 M
┌ Warning: `training` is set to `Val{true}()` but is not being used within an autodiff call (gradient, jacobian, etc...). This will be slow. If you are using a `Lux.jl` model, set it to inference (test) mode using `LuxCore.testmode`. Reliance on this behavior is discouraged, and is not guaranteed by Semantic Versioning, and might be removed without a deprecation cycle. It is recommended to fix this issue in your code.
└ @ LuxLib.Utils /var/lib/buildkite-agent/builds/gpuci-15/julialang/lux-dot-jl/lib/LuxLib/src/utils.jl:334
Epoch 1, Iter 39, Loss: 24262.5761719, Throughput: 50.058082 im/s
Epoch 1, Train Loss: 39666.1914062, Time: 100.1321s, Throughput: 49.854118 im/s
Epoch 2, Iter 39, Loss: 18327.4843750, Throughput: 2154.884072 im/s
Epoch 2, Train Loss: 20089.2832031, Time: 2.3169s, Throughput: 2154.631056 im/s
Epoch 3, Iter 39, Loss: 15309.5693359, Throughput: 2185.270676 im/s
Epoch 3, Train Loss: 16512.3281250, Time: 2.2846s, Throughput: 2185.047870 im/s
Epoch 4, Iter 39, Loss: 14522.6289062, Throughput: 2155.200371 im/s
Epoch 4, Train Loss: 14950.2119141, Time: 2.3166s, Throughput: 2154.887843 im/s
Epoch 5, Iter 39, Loss: 14176.5859375, Throughput: 2185.100546 im/s
Epoch 5, Train Loss: 13952.6220703, Time: 2.2849s, Throughput: 2184.809836 im/s
Epoch 6, Iter 39, Loss: 13097.6044922, Throughput: 2181.017406 im/s
Epoch 6, Train Loss: 13326.8378906, Time: 2.2892s, Throughput: 2180.692577 im/s
Epoch 7, Iter 39, Loss: 12619.8964844, Throughput: 2192.146554 im/s
Epoch 7, Train Loss: 12868.1943359, Time: 2.2777s, Throughput: 2191.725023 im/s
Epoch 8, Iter 39, Loss: 11840.8447266, Throughput: 2184.456071 im/s
Epoch 8, Train Loss: 12519.2626953, Time: 2.2855s, Throughput: 2184.188317 im/s
Epoch 9, Iter 39, Loss: 12160.9150391, Throughput: 2169.659347 im/s
Epoch 9, Train Loss: 12099.7177734, Time: 2.3011s, Throughput: 2169.394533 im/s
Epoch 10, Iter 39, Loss: 12134.8847656, Throughput: 2183.118638 im/s
Epoch 10, Train Loss: 11925.9306641, Time: 2.2869s, Throughput: 2182.871465 im/s
Epoch 11, Iter 39, Loss: 10728.8710938, Throughput: 2194.534697 im/s
Epoch 11, Train Loss: 11774.2255859, Time: 2.2751s, Throughput: 2194.229744 im/s
Epoch 12, Iter 39, Loss: 11998.2080078, Throughput: 2178.963072 im/s
Epoch 12, Train Loss: 11504.6044922, Time: 2.2914s, Throughput: 2178.599637 im/s
Epoch 13, Iter 39, Loss: 11695.4726562, Throughput: 2193.994532 im/s
Epoch 13, Train Loss: 11349.2812500, Time: 2.2755s, Throughput: 2193.766956 im/s
Epoch 14, Iter 39, Loss: 11746.2080078, Throughput: 2189.192365 im/s
Epoch 14, Train Loss: 11280.8144531, Time: 2.2806s, Throughput: 2188.921618 im/s
Epoch 15, Iter 39, Loss: 11914.6582031, Throughput: 2191.547004 im/s
Epoch 15, Train Loss: 11210.9599609, Time: 2.2781s, Throughput: 2191.298837 im/s
Epoch 16, Iter 39, Loss: 10314.3828125, Throughput: 2158.921281 im/s
Epoch 16, Train Loss: 11072.0244141, Time: 2.3125s, Throughput: 2158.700032 im/s
Epoch 17, Iter 39, Loss: 10867.2597656, Throughput: 2165.633881 im/s
Epoch 17, Train Loss: 10934.3847656, Time: 2.3054s, Throughput: 2165.345416 im/s
Epoch 18, Iter 39, Loss: 9800.8828125, Throughput: 2188.217485 im/s
Epoch 18, Train Loss: 10723.6044922, Time: 2.2817s, Throughput: 2187.862616 im/s
Epoch 19, Iter 39, Loss: 10232.0458984, Throughput: 2137.253329 im/s
Epoch 19, Train Loss: 10672.5380859, Time: 2.3361s, Throughput: 2136.933116 im/s
Epoch 20, Iter 39, Loss: 11296.2148438, Throughput: 2189.147274 im/s
Epoch 20, Train Loss: 10596.6103516, Time: 2.2807s, Throughput: 2188.796909 im/s
Epoch 21, Iter 39, Loss: 10113.4746094, Throughput: 2172.620533 im/s
Epoch 21, Train Loss: 10539.1679688, Time: 2.2979s, Throughput: 2172.377760 im/s
Epoch 22, Iter 39, Loss: 10446.4023438, Throughput: 2166.033783 im/s
Epoch 22, Train Loss: 10468.0439453, Time: 2.3051s, Throughput: 2165.659865 im/s
Epoch 23, Iter 39, Loss: 9927.3164062, Throughput: 2177.075802 im/s
Epoch 23, Train Loss: 10330.6064453, Time: 2.2932s, Throughput: 2176.856475 im/s
Epoch 24, Iter 39, Loss: 10710.1728516, Throughput: 2169.761198 im/s
Epoch 24, Train Loss: 10324.8916016, Time: 2.3010s, Throughput: 2169.465114 im/s
Epoch 25, Iter 39, Loss: 10322.6367188, Throughput: 2149.017495 im/s
Epoch 25, Train Loss: 10299.8222656, Time: 2.3236s, Throughput: 2148.400519 im/s
Epoch 26, Iter 39, Loss: 9943.3164062, Throughput: 2182.757911 im/s
Epoch 26, Train Loss: 10223.8027344, Time: 2.2872s, Throughput: 2182.556548 im/s
Epoch 27, Iter 39, Loss: 10085.9736328, Throughput: 2191.732823 im/s
Epoch 27, Train Loss: 10122.6865234, Time: 2.2779s, Throughput: 2191.454795 im/s
Epoch 28, Iter 39, Loss: 10689.1298828, Throughput: 2165.956928 im/s
Epoch 28, Train Loss: 10112.2460938, Time: 2.3050s, Throughput: 2165.693017 im/s
Epoch 29, Iter 39, Loss: 9791.1064453, Throughput: 2182.097534 im/s
Epoch 29, Train Loss: 10042.0273438, Time: 2.2881s, Throughput: 2181.742828 im/s
Epoch 30, Iter 39, Loss: 10382.9824219, Throughput: 2174.263438 im/s
Epoch 30, Train Loss: 10025.1142578, Time: 2.2962s, Throughput: 2174.057769 im/s
Epoch 31, Iter 39, Loss: 10323.5429688, Throughput: 2168.222305 im/s
Epoch 31, Train Loss: 9974.0507812, Time: 2.3027s, Throughput: 2167.928436 im/s
Epoch 32, Iter 39, Loss: 9416.3505859, Throughput: 1920.914625 im/s
Epoch 32, Train Loss: 9984.5869141, Time: 2.5991s, Throughput: 1920.683438 im/s
Epoch 33, Iter 39, Loss: 9853.8574219, Throughput: 2082.496307 im/s
Epoch 33, Train Loss: 9859.9667969, Time: 2.3975s, Throughput: 2082.195811 im/s
Epoch 34, Iter 39, Loss: 10251.0126953, Throughput: 2128.076031 im/s
Epoch 34, Train Loss: 9854.0722656, Time: 2.3461s, Throughput: 2127.774779 im/s
Epoch 35, Iter 39, Loss: 9852.1855469, Throughput: 2167.376833 im/s
Epoch 35, Train Loss: 9815.8505859, Time: 2.3036s, Throughput: 2167.077810 im/s
Epoch 36, Iter 39, Loss: 10304.3984375, Throughput: 2176.081375 im/s
Epoch 36, Train Loss: 9741.7773438, Time: 2.2943s, Throughput: 2175.807303 im/s
Epoch 37, Iter 39, Loss: 9712.1074219, Throughput: 2164.632645 im/s
Epoch 37, Train Loss: 9766.7558594, Time: 2.3065s, Throughput: 2164.275093 im/s
Epoch 38, Iter 39, Loss: 9485.1152344, Throughput: 2165.100236 im/s
Epoch 38, Train Loss: 9715.9580078, Time: 2.3076s, Throughput: 2163.268631 im/s
Epoch 39, Iter 39, Loss: 10443.6943359, Throughput: 2181.658034 im/s
Epoch 39, Train Loss: 9768.0468750, Time: 2.2889s, Throughput: 2180.968107 im/s
Epoch 40, Iter 39, Loss: 9532.7949219, Throughput: 2168.796803 im/s
Epoch 40, Train Loss: 9698.7792969, Time: 2.3020s, Throughput: 2168.525462 im/s
Epoch 41, Iter 39, Loss: 10083.5107422, Throughput: 2153.929087 im/s
Epoch 41, Train Loss: 9653.2880859, Time: 2.3181s, Throughput: 2153.511935 im/s
Epoch 42, Iter 39, Loss: 9803.3671875, Throughput: 2178.284139 im/s
Epoch 42, Train Loss: 9614.5947266, Time: 2.2921s, Throughput: 2177.932711 im/s
Epoch 43, Iter 39, Loss: 9036.3994141, Throughput: 2107.954846 im/s
Epoch 43, Train Loss: 9554.7353516, Time: 2.3685s, Throughput: 2107.641655 im/s
Epoch 44, Iter 39, Loss: 9514.6689453, Throughput: 2149.224851 im/s
Epoch 44, Train Loss: 9487.9042969, Time: 2.3229s, Throughput: 2149.025877 im/s
Epoch 45, Iter 39, Loss: 9114.8886719, Throughput: 2184.505755 im/s
Epoch 45, Train Loss: 9473.5439453, Time: 2.2854s, Throughput: 2184.296550 im/s
Epoch 46, Iter 39, Loss: 9747.8535156, Throughput: 2185.958764 im/s
Epoch 46, Train Loss: 9464.4843750, Time: 2.2839s, Throughput: 2185.766394 im/s
Epoch 47, Iter 39, Loss: 9229.9648438, Throughput: 2176.092683 im/s
Epoch 47, Train Loss: 9403.1728516, Time: 2.2943s, Throughput: 2175.866996 im/s
Epoch 48, Iter 39, Loss: 9158.3251953, Throughput: 2175.128302 im/s
Epoch 48, Train Loss: 9408.1220703, Time: 2.2953s, Throughput: 2174.917952 im/s
Epoch 49, Iter 39, Loss: 8965.3710938, Throughput: 2159.385294 im/s
Epoch 49, Train Loss: 9416.3515625, Time: 2.3122s, Throughput: 2158.967807 im/s
Epoch 50, Iter 39, Loss: 9822.1914062, Throughput: 2190.493024 im/s
Epoch 50, Train Loss: 9380.4169922, Time: 2.2792s, Throughput: 2190.242116 im/s
Appendix
julia
using InteractiveUtils
InteractiveUtils.versioninfo()
if @isdefined(MLDataDevices)
if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice)
println()
CUDA.versioninfo()
end
if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice)
println()
AMDGPU.versioninfo()
end
endJulia Version 1.11.7
Commit f2b3dbda30a (2025-09-08 12:10 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: Linux (x86_64-linux-gnu)
CPU: 48 × AMD EPYC 7402 24-Core Processor
WORD_SIZE: 64
LLVM: libLLVM-16.0.6 (ORCJIT, znver2)
Threads: 48 default, 0 interactive, 24 GC (on 2 virtual cores)
Environment:
JULIA_CPU_THREADS = 2
JULIA_DEPOT_PATH = /root/.cache/julia-buildkite-plugin/depots/01872db4-8c79-43af-ab7d-12abac4f24f6
LD_LIBRARY_PATH = /usr/local/nvidia/lib:/usr/local/nvidia/lib64
JULIA_PKG_SERVER =
JULIA_NUM_THREADS = 48
JULIA_CUDA_HARD_MEMORY_LIMIT = 100%
JULIA_PKG_PRECOMPILE_AUTO = 0
JULIA_DEBUG = LiterateThis page was generated using Literate.jl.