Fitting a Polynomial using MLP

In this tutorial we will fit a MultiLayer Perceptron (MLP) on data generated from a polynomial.

Package Imports

using Lux, ADTypes, LuxCUDA, Optimisers, Printf, Random, Statistics, Zygote
using CairoMakie

Precompiling CairoMakie...
    444.4 ms  ✓ StatsAPI
    456.7 ms  ✓ PCRE2_jll
    554.6 ms  ✓ TranscodingStreams
    507.6 ms  ✓ IntervalSets → IntervalSetsStatisticsExt
    395.7 ms  ✓ Ratios → RatiosFixedPointNumbersExt
    687.8 ms  ✓ SharedArrays
    860.2 ms  ✓ ProgressMeter
    692.5 ms  ✓ WoodburyMatrices
    900.9 ms  ✓ PDMats
    364.4 ms  ✓ ConstructionBase → ConstructionBaseIntervalSetsExt
    786.6 ms  ✓ AxisArrays
   1059.3 ms  ✓ GeoInterface
    407.4 ms  ✓ Showoff
    387.4 ms  ✓ RelocatableFolders
    577.2 ms  ✓ Graphics
    677.2 ms  ✓ Animations
   2218.2 ms  ✓ ColorVectorSpace
    618.8 ms  ✓ InverseFunctions → InverseFunctionsTestExt
    365.7 ms  ✓ SignedDistanceFields
   1573.9 ms  ✓ UnicodeFun
    643.9 ms  ✓ OpenSSL_jll
    615.3 ms  ✓ Graphite2_jll
    699.0 ms  ✓ Libmount_jll
    663.7 ms  ✓ Bzip2_jll
    601.5 ms  ✓ Rmath_jll
    583.2 ms  ✓ Xorg_libXau_jll
    607.2 ms  ✓ libpng_jll
    582.1 ms  ✓ libfdk_aac_jll
    583.4 ms  ✓ Imath_jll
    576.2 ms  ✓ Giflib_jll
    583.1 ms  ✓ LAME_jll
    607.3 ms  ✓ Pixman_jll
    572.5 ms  ✓ LERC_jll
    567.5 ms  ✓ EarCut_jll
    572.1 ms  ✓ CRlibm_jll
    626.3 ms  ✓ JpegTurbo_jll
    651.1 ms  ✓ XZ_jll
    602.3 ms  ✓ Ogg_jll
    579.6 ms  ✓ Xorg_libXdmcp_jll
    595.8 ms  ✓ x265_jll
    583.2 ms  ✓ x264_jll
    615.9 ms  ✓ libaom_jll
    650.5 ms  ✓ Zstd_jll
    633.0 ms  ✓ Expat_jll
    605.3 ms  ✓ LZO_jll
    579.3 ms  ✓ Opus_jll
    497.9 ms  ✓ Xorg_xtrans_jll
    583.6 ms  ✓ Libffi_jll
    590.7 ms  ✓ Libgpg_error_jll
    543.2 ms  ✓ isoband_jll
    606.2 ms  ✓ FFTW_jll
    549.4 ms  ✓ Xorg_libpthread_stubs_jll
    563.5 ms  ✓ Libuuid_jll
    599.7 ms  ✓ FriBidi_jll
   1049.6 ms  ✓ FilePathsBase
    477.9 ms  ✓ LogExpFunctions → LogExpFunctionsInverseFunctionsExt
   4448.0 ms  ✓ FileIO
    675.8 ms  ✓ XML2_jll
   1008.0 ms  ✓ QuadGK
  22244.6 ms  ✓ Unitful
   8713.8 ms  ✓ SIMD
   1152.8 ms  ✓ HypergeometricFunctions
   1253.9 ms  ✓ ColorBrewer
    658.4 ms  ✓ AxisAlgorithms
    666.4 ms  ✓ FillArrays → FillArraysPDMatsExt
    736.1 ms  ✓ ColorVectorSpace → SpecialFunctionsExt
   2333.2 ms  ✓ StatsBase
    621.7 ms  ✓ FreeType2_jll
    843.9 ms  ✓ Rmath
    755.5 ms  ✓ OpenEXR_jll
   3522.7 ms  ✓ ColorSchemes
   2536.2 ms  ✓ IntervalArithmetic
    691.9 ms  ✓ libsixel_jll
    674.9 ms  ✓ libvorbis_jll
    673.3 ms  ✓ Libtiff_jll
   5844.9 ms  ✓ GeometryBasics
    640.2 ms  ✓ Libgcrypt_jll
    478.7 ms  ✓ Isoband
    552.2 ms  ✓ FilePathsBase → FilePathsBaseMmapExt
   1213.0 ms  ✓ FilePathsBase → FilePathsBaseTestExt
    812.0 ms  ✓ FilePaths
   1539.8 ms  ✓ QOI
    713.6 ms  ✓ Gettext_jll
    585.8 ms  ✓ Unitful → ConstructionBaseUnitfulExt
    579.4 ms  ✓ Unitful → InverseFunctionsUnitfulExt
  10126.7 ms  ✓ FFTW
  20030.4 ms  ✓ ImageCore
   2065.6 ms  ✓ Interpolations
   1007.8 ms  ✓ FreeType
    817.0 ms  ✓ Fontconfig_jll
   1926.7 ms  ✓ StatsFuns
   1598.4 ms  ✓ OpenEXR
  13551.5 ms  ✓ Automa
   4662.0 ms  ✓ ExactPredicates
    534.9 ms  ✓ IntervalArithmetic → IntervalArithmeticIntervalSetsExt
  11935.6 ms  ✓ PlotUtils
   1222.9 ms  ✓ Packing
   6383.6 ms  ✓ GridLayoutBase
   1381.1 ms  ✓ ShaderAbstractions
    713.2 ms  ✓ XSLT_jll
    832.4 ms  ✓ Glib_jll
   2233.3 ms  ✓ ImageBase
   4679.8 ms  ✓ MakieCore
   3221.2 ms  ✓ JpegTurbo
   4161.2 ms  ✓ PNGFiles
   1190.3 ms  ✓ Interpolations → InterpolationsUnitfulExt
   2205.7 ms  ✓ FreeTypeAbstraction
   4812.9 ms  ✓ Sixel
    675.4 ms  ✓ StatsFuns → StatsFunsInverseFunctionsExt
   1562.8 ms  ✓ StatsFuns → StatsFunsChainRulesCoreExt
   5050.4 ms  ✓ Distributions
   1165.5 ms  ✓ Xorg_libxcb_jll
   5600.3 ms  ✓ DelaunayTriangulation
   2131.9 ms  ✓ ImageAxes
   1536.7 ms  ✓ Distributions → DistributionsTestExt
   1585.9 ms  ✓ Distributions → DistributionsChainRulesCoreExt
    634.1 ms  ✓ Xorg_libX11_jll
   1138.4 ms  ✓ ImageMetadata
   1835.2 ms  ✓ KernelDensity
    649.4 ms  ✓ Xorg_libXrender_jll
    695.5 ms  ✓ Xorg_libXext_jll
   2024.2 ms  ✓ Netpbm
    870.9 ms  ✓ Cairo_jll
    780.3 ms  ✓ Libglvnd_jll
    811.6 ms  ✓ HarfBuzz_jll
    876.9 ms  ✓ libwebp_jll
  15166.6 ms  ✓ MathTeXEngine
    804.9 ms  ✓ libass_jll
    866.3 ms  ✓ Pango_jll
    940.1 ms  ✓ FFMPEG_jll
   2690.6 ms  ✓ WebP
   1340.7 ms  ✓ Cairo
  67080.3 ms  ✓ TiffImages
   1245.0 ms  ✓ ImageIO
 161372.9 ms  ✓ Makie
  92993.1 ms  ✓ CairoMakie
  136 dependencies successfully precompiled in 370 seconds. 133 already precompiled.
Precompiling NNlibFFTWExt...
    924.5 ms  ✓ NNlib → NNlibFFTWExt
  1 dependency successfully precompiled in 1 seconds. 54 already precompiled.
Precompiling IntervalArithmeticForwardDiffExt...
    491.3 ms  ✓ IntervalArithmetic → IntervalArithmeticDiffRulesExt
    702.8 ms  ✓ IntervalArithmetic → IntervalArithmeticForwardDiffExt
  2 dependencies successfully precompiled in 1 seconds. 42 already precompiled.

Dataset

Generate 128 datapoints from the polynomial $y=x^2-2x$.

function generate_data(rng::AbstractRNG)
    x = reshape(collect(range(-2.0f0, 2.0f0, 128)), (1, 128))
    y = evalpoly.(x, ((0, -2, 1),)) .+ randn(rng, Float32, (1, 128)) .* 0.1f0
    return (x, y)
end

generate_data (generic function with 1 method)

Initialize the random number generator and fetch the dataset.

rng = MersenneTwister()
Random.seed!(rng, 12345)

(x, y) = generate_data(rng)

(Float32[-2.0 -1.968504 -1.9370079 -1.9055119 -1.8740157 -1.8425196 -1.8110236 -1.7795275 -1.7480315 -1.7165354 -1.6850394 -1.6535434 -1.6220472 -1.5905511 -1.5590551 -1.527559 -1.496063 -1.464567 -1.4330709 -1.4015749 -1.3700787 -1.3385826 -1.3070866 -1.2755905 -1.2440945 -1.2125984 -1.1811024 -1.1496063 -1.1181102 -1.0866141 -1.0551181 -1.023622 -0.992126 -0.96062994 -0.92913383 -0.8976378 -0.86614174 -0.8346457 -0.8031496 -0.77165353 -0.7401575 -0.70866144 -0.6771653 -0.6456693 -0.61417323 -0.5826772 -0.5511811 -0.51968503 -0.48818898 -0.4566929 -0.42519686 -0.39370078 -0.36220473 -0.33070865 -0.2992126 -0.26771653 -0.23622048 -0.20472442 -0.17322835 -0.14173229 -0.11023622 -0.07874016 -0.047244094 -0.015748031 0.015748031 0.047244094 0.07874016 0.11023622 0.14173229 0.17322835 0.20472442 0.23622048 0.26771653 0.2992126 0.33070865 0.36220473 0.39370078 0.42519686 0.4566929 0.48818898 0.51968503 0.5511811 0.5826772 0.61417323 0.6456693 0.6771653 0.70866144 0.7401575 0.77165353 0.8031496 0.8346457 0.86614174 0.8976378 0.92913383 0.96062994 0.992126 1.023622 1.0551181 1.0866141 1.1181102 1.1496063 1.1811024 1.2125984 1.2440945 1.2755905 1.3070866 1.3385826 1.3700787 1.4015749 1.4330709 1.464567 1.496063 1.527559 1.5590551 1.5905511 1.6220472 1.6535434 1.6850394 1.7165354 1.7480315 1.7795275 1.8110236 1.8425196 1.8740157 1.9055119 1.9370079 1.968504 2.0], Float32[8.080871 7.562357 7.451749 7.5005703 7.295229 7.2245107 6.8731666 6.7092047 6.5385857 6.4631066 6.281978 5.960991 5.963052 5.68927 5.3667717 5.519665 5.2999034 5.0238676 5.174298 4.6706038 4.570324 4.439068 4.4462147 4.299262 3.9799082 3.9492173 3.8747025 3.7264304 3.3844414 3.2934628 3.1180353 3.0698316 3.0491123 2.592982 2.8164148 2.3875027 2.3781595 2.4269633 2.2763796 2.3316176 2.0829067 1.9049499 1.8581494 1.7632381 1.7745113 1.5406592 1.3689325 1.2614254 1.1482575 1.2801026 0.9070533 0.91188717 0.9415703 0.85747254 0.6692604 0.7172643 0.48259094 0.48990166 0.35299227 0.31578436 0.25483933 0.37486005 0.19847682 -0.042415008 -0.05951088 0.014774345 -0.114184186 -0.15978265 -0.29916334 -0.22005874 -0.17161606 -0.3613516 -0.5489093 -0.7267406 -0.5943626 -0.62129945 -0.50063384 -0.6346849 -0.86081326 -0.58715504 -0.5171875 -0.6575044 -0.71243864 -0.78395927 -0.90537953 -0.9515314 -0.8603811 -0.92880917 -1.0078154 -0.90215015 -1.0109437 -1.0764086 -1.1691734 -1.0740278 -1.1429857 -1.104191 -0.948015 -0.9233653 -0.82379496 -0.9810639 -0.92863405 -0.9360056 -0.92652786 -0.847396 -1.115507 -1.0877254 -0.92295444 -0.86975616 -0.81879705 -0.8482455 -0.6524158 -0.6184501 -0.7483137 -0.60395515 -0.67555165 -0.6288941 -0.6774449 -0.49889082 -0.43817532 -0.46497717 -0.30316323 -0.36745527 -0.3227286 -0.20977046 -0.09777648 -0.053120755 -0.15877295 -0.06777584])

Let's visualize the dataset

begin
    fig = Figure()
    ax = CairoMakie.Axis(fig[1, 1]; xlabel="x", ylabel="y")

    l = lines!(ax, x[1, :], x -> evalpoly(x, (0, -2, 1)); linewidth=3, color=:blue)
    s = scatter!(ax, x[1, :], y[1, :]; markersize=12, alpha=0.5,
        color=:orange, strokecolor=:black, strokewidth=2)

    axislegend(ax, [l, s], ["True Quadratic Function", "Data Points"])

    fig
end

Neural Network

For this problem, you should not be using a neural network. But let's still do that!

model = Chain(Dense(1 => 16, relu), Dense(16 => 1))

Chain(
    layer_1 = Dense(1 => 16, relu),     # 32 parameters
    layer_2 = Dense(16 => 1),           # 17 parameters
)         # Total: 49 parameters,
          #        plus 0 states.

Optimizer

We will use Adam from Optimisers.jl

opt = Adam(0.03f0)

Adam(0.03, (0.9, 0.999), 1.0e-8)

Loss Function

We will use the Training API so we need to ensure that our loss function takes 4 inputs – model, parameters, states and data. The function must return 3 values – loss, updated_state, and any computed statistics. This is already satisfied by the loss functions provided by Lux.

const loss_function = MSELoss()

const dev_cpu = cpu_device()
const dev_gpu = gpu_device()

ps, st = Lux.setup(rng, model) |> dev_gpu

((layer_1 = (weight = Float32[2.2569513; 1.8385266; 1.8834435; -1.4215803; -0.1289033; -1.4116536; -1.4359436; -2.3610642; -0.847535; 1.6091344; -0.34999675; 1.9372884; -0.41628727; 1.1786895; -1.4312565; 0.34652048;;], bias = Float32[0.9155488, -0.005158901, 0.5026965, -0.84174657, -0.9167142, -0.14881086, -0.8202727, 0.19286752, 0.60171676, 0.951689, 0.4595859, -0.33281517, -0.692657, 0.4369135, 0.3800323, 0.61768365]), layer_2 = (weight = Float32[0.20061705 0.22529833 0.07667785 0.115506485 0.22827768 0.22680467 0.0035893882 -0.39495495 0.18033011 -0.02850357 -0.08613788 -0.3103005 0.12508307 -0.087390475 -0.13759731 0.08034529], bias = Float32[0.06066203])), (layer_1 = NamedTuple(), layer_2 = NamedTuple()))

Training

First we will create a Training.TrainState which is essentially a convenience wrapper over parameters, states and optimizer states.

tstate = Training.TrainState(model, ps, st, opt)

TrainState
    model: Lux.Chain{@NamedTuple{layer_1::Lux.Dense{typeof(NNlib.relu), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Lux.Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = Dense(1 => 16, relu), layer_2 = Dense(16 => 1)), nothing)
    # of parameters: 49
    # of states: 0
    optimizer: Adam(0.03, (0.9, 0.999), 1.0e-8)
    step: 0

Now we will use Zygote for our AD requirements.

vjp_rule = AutoZygote()

ADTypes.AutoZygote()

Finally the training loop.

function main(tstate::Training.TrainState, vjp, data, epochs)
    data = data .|> gpu_device()
    for epoch in 1:epochs
        _, loss, _, tstate = Training.single_train_step!(vjp, loss_function, data, tstate)
        if epoch % 50 == 1 || epoch == epochs
            @printf "Epoch: %3d \t Loss: %.5g\n" epoch loss
        end
    end
    return tstate
end

tstate = main(tstate, vjp_rule, (x, y), 250)
y_pred = dev_cpu(Lux.apply(tstate.model, dev_gpu(x), tstate.parameters, tstate.states)[1])

Epoch:   1 	 Loss: 12.499
Epoch:  51 	 Loss: 0.090343
Epoch: 101 	 Loss: 0.0405
Epoch: 151 	 Loss: 0.024803
Epoch: 201 	 Loss: 0.017711
Epoch: 250 	 Loss: 0.01466

Let's plot the results

begin
    fig = Figure()
    ax = CairoMakie.Axis(fig[1, 1]; xlabel="x", ylabel="y")

    l = lines!(ax, x[1, :], x -> evalpoly(x, (0, -2, 1)); linewidth=3)
    s1 = scatter!(ax, x[1, :], y[1, :]; markersize=12, alpha=0.5,
        color=:orange, strokecolor=:black, strokewidth=2)
    s2 = scatter!(ax, x[1, :], y_pred[1, :]; markersize=12, alpha=0.5,
        color=:green, strokecolor=:black, strokewidth=2)

    axislegend(ax, [l, s1, s2], ["True Quadratic Function", "Actual Data", "Predictions"])

    fig
end

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.11.2
Commit 5e9a32e7af2 (2024-12-01 20:02 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 = Literate

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

CUDA libraries: 
- CUBLAS: 12.6.4
- 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.4+0
- CUDA_Runtime_jll: 0.15.5+0

Toolchain:
- Julia: 1.11.2
- LLVM: 16.0.6

Environment:
- JULIA_CUDA_HARD_MEMORY_LIMIT: 100%

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

---

This page was generated using Literate.jl.