Training a HyperNetwork on MNIST and FashionMNIST

Package Imports

using Lux, ADTypes, ComponentArrays, LuxCUDA, MLDatasets, MLUtils, OneHotArrays, Optimisers,
      Printf, Random, Setfield, Statistics, Zygote

CUDA.allowscalar(false)

Loading Datasets

function load_dataset(::Type{dset}, n_train::Int, n_eval::Int, batchsize::Int) where {dset}
    imgs, labels = dset(:train)[1:n_train]
    x_train, y_train = reshape(imgs, 28, 28, 1, n_train), onehotbatch(labels, 0:9)

    imgs, labels = dset(:test)[1:n_eval]
    x_test, y_test = reshape(imgs, 28, 28, 1, n_eval), onehotbatch(labels, 0:9)

    return (
        DataLoader((x_train, y_train); batchsize=min(batchsize, n_train), shuffle=true),
        DataLoader((x_test, y_test); batchsize=min(batchsize, n_eval), shuffle=false)
    )
end

function load_datasets(n_train=1024, n_eval=32, batchsize=256)
    return load_dataset.((MNIST, FashionMNIST), n_train, n_eval, batchsize)
end

load_datasets (generic function with 4 methods)

Implement a HyperNet Layer

function HyperNet(
        weight_generator::Lux.AbstractLuxLayer, core_network::Lux.AbstractLuxLayer)
    ca_axes = Lux.initialparameters(Random.default_rng(), core_network) |>
              ComponentArray |>
              getaxes
    return @compact(; ca_axes, weight_generator, core_network, dispatch=:HyperNet) do (x, y)
        # Generate the weights
        ps_new = ComponentArray(vec(weight_generator(x)), ca_axes)
        @return core_network(y, ps_new)
    end
end

HyperNet (generic function with 1 method)

Defining functions on the CompactLuxLayer requires some understanding of how the layer is structured, as such we don't recommend doing it unless you are familiar with the internals. In this case, we simply write it to ignore the initialization of the core_network parameters.

function Lux.initialparameters(rng::AbstractRNG, hn::CompactLuxLayer{:HyperNet})
    return (; weight_generator=Lux.initialparameters(rng, hn.layers.weight_generator),)
end

Create and Initialize the HyperNet

function create_model()
    # Doesn't need to be a MLP can have any Lux Layer
    core_network = Chain(FlattenLayer(), Dense(784, 256, relu), Dense(256, 10))
    weight_generator = Chain(
        Embedding(2 => 32),
        Dense(32, 64, relu),
        Dense(64, Lux.parameterlength(core_network))
    )

    model = HyperNet(weight_generator, core_network)
    return model
end

create_model (generic function with 1 method)

Define Utility Functions

const loss = CrossEntropyLoss(; logits=Val(true))

function accuracy(model, ps, st, dataloader, data_idx)
    total_correct, total = 0, 0
    st = Lux.testmode(st)
    for (x, y) in dataloader
        target_class = onecold(y)
        predicted_class = onecold(first(model((data_idx, x), ps, st)))
        total_correct += sum(target_class .== predicted_class)
        total += length(target_class)
    end
    return total_correct / total
end

accuracy (generic function with 1 method)

Training

function train()
    model = create_model()
    dataloaders = load_datasets()

    dev = gpu_device()
    rng = Xoshiro(0)
    ps, st = Lux.setup(rng, model) |> dev

    train_state = Training.TrainState(model, ps, st, Adam(0.001f0))

    ### Lets train the model
    nepochs = 50
    for epoch in 1:nepochs, data_idx in 1:2
        train_dataloader, test_dataloader = dataloaders[data_idx] .|> dev

        stime = time()
        for (x, y) in train_dataloader
            (_, _, _, train_state) = Training.single_train_step!(
                AutoZygote(), loss, ((data_idx, x), y), train_state)
        end
        ttime = time() - stime

        train_acc = round(
            accuracy(model, train_state.parameters,
                train_state.states, train_dataloader, data_idx) * 100;
            digits=2)
        test_acc = round(
            accuracy(model, train_state.parameters,
                train_state.states, test_dataloader, data_idx) * 100;
            digits=2)

        data_name = data_idx == 1 ? "MNIST" : "FashionMNIST"

        @printf "[%3d/%3d]\t%12s\tTime %3.5fs\tTraining Accuracy: %3.2f%%\tTest \
                 Accuracy: %3.2f%%\n" epoch nepochs data_name ttime train_acc test_acc
    end

    println()

    test_acc_list = [0.0, 0.0]
    for data_idx in 1:2
        train_dataloader, test_dataloader = dataloaders[data_idx] .|> dev
        train_acc = round(
            accuracy(model, train_state.parameters,
                train_state.states, train_dataloader, data_idx) * 100;
            digits=2)
        test_acc = round(
            accuracy(model, train_state.parameters,
                train_state.states, test_dataloader, data_idx) * 100;
            digits=2)

        data_name = data_idx == 1 ? "MNIST" : "FashionMNIST"

        @printf "[FINAL]\t%12s\tTraining Accuracy: %3.2f%%\tTest Accuracy: \
                 %3.2f%%\n" data_name train_acc test_acc
        test_acc_list[data_idx] = test_acc
    end
    return test_acc_list
end

test_acc_list = train()

[  1/ 50]	       MNIST	Time 72.54884s	Training Accuracy: 58.59%	Test Accuracy: 56.25%
[  1/ 50]	FashionMNIST	Time 0.02819s	Training Accuracy: 47.17%	Test Accuracy: 37.50%
[  2/ 50]	       MNIST	Time 0.02820s	Training Accuracy: 67.29%	Test Accuracy: 59.38%
[  2/ 50]	FashionMNIST	Time 0.02929s	Training Accuracy: 59.86%	Test Accuracy: 56.25%
[  3/ 50]	       MNIST	Time 0.02961s	Training Accuracy: 75.98%	Test Accuracy: 68.75%
[  3/ 50]	FashionMNIST	Time 0.02841s	Training Accuracy: 65.43%	Test Accuracy: 53.12%
[  4/ 50]	       MNIST	Time 0.02648s	Training Accuracy: 77.54%	Test Accuracy: 62.50%
[  4/ 50]	FashionMNIST	Time 0.02292s	Training Accuracy: 68.26%	Test Accuracy: 65.62%
[  5/ 50]	       MNIST	Time 0.02207s	Training Accuracy: 83.59%	Test Accuracy: 65.62%
[  5/ 50]	FashionMNIST	Time 0.06605s	Training Accuracy: 70.70%	Test Accuracy: 56.25%
[  6/ 50]	       MNIST	Time 0.02071s	Training Accuracy: 86.52%	Test Accuracy: 62.50%
[  6/ 50]	FashionMNIST	Time 0.02103s	Training Accuracy: 74.41%	Test Accuracy: 62.50%
[  7/ 50]	       MNIST	Time 0.02099s	Training Accuracy: 87.60%	Test Accuracy: 71.88%
[  7/ 50]	FashionMNIST	Time 0.02107s	Training Accuracy: 76.07%	Test Accuracy: 71.88%
[  8/ 50]	       MNIST	Time 0.02105s	Training Accuracy: 91.31%	Test Accuracy: 71.88%
[  8/ 50]	FashionMNIST	Time 0.02019s	Training Accuracy: 80.37%	Test Accuracy: 75.00%
[  9/ 50]	       MNIST	Time 0.02098s	Training Accuracy: 93.46%	Test Accuracy: 71.88%
[  9/ 50]	FashionMNIST	Time 0.02080s	Training Accuracy: 82.13%	Test Accuracy: 78.12%
[ 10/ 50]	       MNIST	Time 0.03527s	Training Accuracy: 95.41%	Test Accuracy: 78.12%
[ 10/ 50]	FashionMNIST	Time 0.02118s	Training Accuracy: 82.13%	Test Accuracy: 71.88%
[ 11/ 50]	       MNIST	Time 0.02078s	Training Accuracy: 97.56%	Test Accuracy: 78.12%
[ 11/ 50]	FashionMNIST	Time 0.02039s	Training Accuracy: 83.69%	Test Accuracy: 78.12%
[ 12/ 50]	       MNIST	Time 0.02071s	Training Accuracy: 97.75%	Test Accuracy: 81.25%
[ 12/ 50]	FashionMNIST	Time 0.02083s	Training Accuracy: 85.55%	Test Accuracy: 78.12%
[ 13/ 50]	       MNIST	Time 0.02105s	Training Accuracy: 98.44%	Test Accuracy: 78.12%
[ 13/ 50]	FashionMNIST	Time 0.02134s	Training Accuracy: 85.25%	Test Accuracy: 78.12%
[ 14/ 50]	       MNIST	Time 0.02049s	Training Accuracy: 98.93%	Test Accuracy: 81.25%
[ 14/ 50]	FashionMNIST	Time 0.02778s	Training Accuracy: 86.91%	Test Accuracy: 84.38%
[ 15/ 50]	       MNIST	Time 0.02143s	Training Accuracy: 99.12%	Test Accuracy: 78.12%
[ 15/ 50]	FashionMNIST	Time 0.02116s	Training Accuracy: 87.40%	Test Accuracy: 84.38%
[ 16/ 50]	       MNIST	Time 0.02102s	Training Accuracy: 99.41%	Test Accuracy: 81.25%
[ 16/ 50]	FashionMNIST	Time 0.02134s	Training Accuracy: 86.23%	Test Accuracy: 78.12%
[ 17/ 50]	       MNIST	Time 0.02142s	Training Accuracy: 99.61%	Test Accuracy: 81.25%
[ 17/ 50]	FashionMNIST	Time 0.02228s	Training Accuracy: 88.28%	Test Accuracy: 84.38%
[ 18/ 50]	       MNIST	Time 0.02262s	Training Accuracy: 99.80%	Test Accuracy: 78.12%
[ 18/ 50]	FashionMNIST	Time 0.02127s	Training Accuracy: 88.96%	Test Accuracy: 75.00%
[ 19/ 50]	       MNIST	Time 0.02164s	Training Accuracy: 99.80%	Test Accuracy: 78.12%
[ 19/ 50]	FashionMNIST	Time 0.02106s	Training Accuracy: 89.26%	Test Accuracy: 81.25%
[ 20/ 50]	       MNIST	Time 0.02094s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 20/ 50]	FashionMNIST	Time 0.02128s	Training Accuracy: 90.72%	Test Accuracy: 84.38%
[ 21/ 50]	       MNIST	Time 0.02068s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 21/ 50]	FashionMNIST	Time 0.02084s	Training Accuracy: 90.92%	Test Accuracy: 84.38%
[ 22/ 50]	       MNIST	Time 0.02110s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 22/ 50]	FashionMNIST	Time 0.02135s	Training Accuracy: 91.31%	Test Accuracy: 84.38%
[ 23/ 50]	       MNIST	Time 0.02848s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 23/ 50]	FashionMNIST	Time 0.02107s	Training Accuracy: 91.60%	Test Accuracy: 84.38%
[ 24/ 50]	       MNIST	Time 0.02073s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 24/ 50]	FashionMNIST	Time 0.02061s	Training Accuracy: 92.09%	Test Accuracy: 84.38%
[ 25/ 50]	       MNIST	Time 0.02081s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 25/ 50]	FashionMNIST	Time 0.02091s	Training Accuracy: 92.19%	Test Accuracy: 84.38%
[ 26/ 50]	       MNIST	Time 0.02090s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 26/ 50]	FashionMNIST	Time 0.02090s	Training Accuracy: 92.97%	Test Accuracy: 84.38%
[ 27/ 50]	       MNIST	Time 0.02225s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 27/ 50]	FashionMNIST	Time 0.02888s	Training Accuracy: 93.55%	Test Accuracy: 84.38%
[ 28/ 50]	       MNIST	Time 0.02092s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 28/ 50]	FashionMNIST	Time 0.02066s	Training Accuracy: 93.85%	Test Accuracy: 84.38%
[ 29/ 50]	       MNIST	Time 0.02065s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 29/ 50]	FashionMNIST	Time 0.02095s	Training Accuracy: 94.34%	Test Accuracy: 84.38%
[ 30/ 50]	       MNIST	Time 0.02080s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 30/ 50]	FashionMNIST	Time 0.02105s	Training Accuracy: 94.34%	Test Accuracy: 84.38%
[ 31/ 50]	       MNIST	Time 0.02091s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 31/ 50]	FashionMNIST	Time 0.02080s	Training Accuracy: 95.02%	Test Accuracy: 84.38%
[ 32/ 50]	       MNIST	Time 0.02714s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 32/ 50]	FashionMNIST	Time 0.02092s	Training Accuracy: 95.31%	Test Accuracy: 84.38%
[ 33/ 50]	       MNIST	Time 0.02036s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 33/ 50]	FashionMNIST	Time 0.02110s	Training Accuracy: 94.92%	Test Accuracy: 84.38%
[ 34/ 50]	       MNIST	Time 0.02097s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 34/ 50]	FashionMNIST	Time 0.02049s	Training Accuracy: 95.41%	Test Accuracy: 84.38%
[ 35/ 50]	       MNIST	Time 0.02102s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 35/ 50]	FashionMNIST	Time 0.02043s	Training Accuracy: 95.90%	Test Accuracy: 84.38%
[ 36/ 50]	       MNIST	Time 0.02131s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 36/ 50]	FashionMNIST	Time 0.02061s	Training Accuracy: 95.12%	Test Accuracy: 84.38%
[ 37/ 50]	       MNIST	Time 0.02070s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 37/ 50]	FashionMNIST	Time 0.02119s	Training Accuracy: 95.21%	Test Accuracy: 84.38%
[ 38/ 50]	       MNIST	Time 0.02067s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 38/ 50]	FashionMNIST	Time 0.02097s	Training Accuracy: 96.19%	Test Accuracy: 84.38%
[ 39/ 50]	       MNIST	Time 0.02092s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 39/ 50]	FashionMNIST	Time 0.02066s	Training Accuracy: 96.19%	Test Accuracy: 84.38%
[ 40/ 50]	       MNIST	Time 0.02040s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 40/ 50]	FashionMNIST	Time 0.02714s	Training Accuracy: 95.41%	Test Accuracy: 81.25%
[ 41/ 50]	       MNIST	Time 0.02037s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 41/ 50]	FashionMNIST	Time 0.02091s	Training Accuracy: 96.09%	Test Accuracy: 84.38%
[ 42/ 50]	       MNIST	Time 0.02127s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 42/ 50]	FashionMNIST	Time 0.02051s	Training Accuracy: 96.88%	Test Accuracy: 84.38%
[ 43/ 50]	       MNIST	Time 0.02094s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 43/ 50]	FashionMNIST	Time 0.02104s	Training Accuracy: 96.97%	Test Accuracy: 84.38%
[ 44/ 50]	       MNIST	Time 0.02108s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 44/ 50]	FashionMNIST	Time 0.02102s	Training Accuracy: 97.27%	Test Accuracy: 84.38%
[ 45/ 50]	       MNIST	Time 0.02515s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 45/ 50]	FashionMNIST	Time 0.02105s	Training Accuracy: 97.36%	Test Accuracy: 84.38%
[ 46/ 50]	       MNIST	Time 0.02119s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 46/ 50]	FashionMNIST	Time 0.02084s	Training Accuracy: 97.95%	Test Accuracy: 84.38%
[ 47/ 50]	       MNIST	Time 0.02115s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 47/ 50]	FashionMNIST	Time 0.02078s	Training Accuracy: 97.75%	Test Accuracy: 84.38%
[ 48/ 50]	       MNIST	Time 0.02118s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 48/ 50]	FashionMNIST	Time 0.02057s	Training Accuracy: 98.05%	Test Accuracy: 84.38%
[ 49/ 50]	       MNIST	Time 0.02071s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 49/ 50]	FashionMNIST	Time 0.02737s	Training Accuracy: 97.66%	Test Accuracy: 84.38%
[ 50/ 50]	       MNIST	Time 0.02098s	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[ 50/ 50]	FashionMNIST	Time 0.02096s	Training Accuracy: 98.14%	Test Accuracy: 84.38%

[FINAL]	       MNIST	Training Accuracy: 100.00%	Test Accuracy: 78.12%
[FINAL]	FashionMNIST	Training Accuracy: 98.14%	Test Accuracy: 84.38%

Appendix

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
end

Julia Version 1.10.6
Commit 67dffc4a8ae (2024-10-28 12:23 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
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (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 = Literate

CUDA runtime 12.6, artifact installation
CUDA driver 12.6
NVIDIA driver 560.35.3

CUDA libraries: 
- CUBLAS: 12.6.3
- CURAND: 10.3.7
- CUFFT: 11.3.0
- CUSOLVER: 11.7.1
- CUSPARSE: 12.5.4
- CUPTI: 2024.3.2 (API 24.0.0)
- NVML: 12.0.0+560.35.3

Julia packages: 
- CUDA: 5.5.2
- CUDA_Driver_jll: 0.10.3+0
- CUDA_Runtime_jll: 0.15.4+0

Toolchain:
- Julia: 1.10.6
- LLVM: 15.0.7

Environment:
- JULIA_CUDA_HARD_MEMORY_LIMIT: 100%

1 device:
  0: NVIDIA A100-PCIE-40GB MIG 1g.5gb (sm_80, 2.170 GiB / 4.750 GiB available)

---

This page was generated using Literate.jl.