diff --git a/.buildkite/testing.yml b/.buildkite/testing.yml index 95e0a68d..3f756e3f 100644 --- a/.buildkite/testing.yml +++ b/.buildkite/testing.yml @@ -1,12 +1,12 @@ steps: - group: ":julia: CUDA GPU" steps: - - label: ":julia: Julia {{matrix.julia}} + CUDA GPU + {{matrix.group}} + Reactant: {{matrix.reactant}}" + - label: ":julia: Julia {{matrix.julia}} + CUDA GPU + {{matrix.group}}" plugins: - JuliaCI/julia#v1: version: "{{matrix.julia}}" - JuliaCI/julia-test#v1: - test_args: "BACKEND_GROUP=CUDA BOLTZ_TEST_GROUP={{matrix.group}} BOLTZ_TEST_REACTANT={{matrix.reactant}}" + test_args: "{{matrix.group}}" - JuliaCI/julia-coverage#v1: codecov: true dirs: @@ -22,19 +22,19 @@ steps: julia: - "1" group: - - "all" - reactant: - - "true" - - "false" + - "layers" + - "vision" + - "piml" + - "core" # - group: ":julia: AMD GPU" # steps: - # - label: ":julia: Julia: {{matrix.julia}} + AMD GPU + {{matrix.group}} + Reactant: {{matrix.reactant}}" + # - label: ":julia: Julia: {{matrix.julia}} + AMD GPU + {{matrix.group}}" # plugins: # - JuliaCI/julia#v1: # version: "{{matrix.julia}}" # - JuliaCI/julia-test#v1: - # test_args: "BACKEND_GROUP=AMDGPU BOLTZ_TEST_GROUP={{matrix.group}} BOLTZ_TEST_REACTANT={{matrix.reactant}}" + # test_args: "{{matrix.group}}" # - JuliaCI/julia-coverage#v1: # codecov: true # dirs: @@ -55,9 +55,10 @@ steps: # julia: # - "1" # group: - # - "all" - # reactant: - # - "false" + # - "layers" + # - "vision" + # - "piml" + # - "core" env: SECRET_CODECOV_TOKEN: "gZlC/IAmeUJehhP5mP2QuUV5a1qV61cvo4PUCLkA9vVkt3x6wgD6fTZmCm+f+gHkmkssFxX+q2h1Ud00XXc75H2LrjyR/cDTIthcO46BBOidYocv/U0gfhp6uT2IZ9fi+ryFfTVVpZ0RIUGmDTj0O/b5qt4oaTriAArLAq6mMipbIR9YCz7ZD/hWQXx8oDeAbnDpwQaddwPyhJhz95nayknOpuJj+ClaVOxgsLGZc3ZWiTj1QxkXBNwxLD2ALeG16Qxs9h7eK87sdcbWeTihvJ6OooARgpoVJAa2pJCFYOGy4Bh07c0VTZmicN2M3GIi74Y5T1PWNaz7nGeANO5Pow==;U2FsdGVkX1843DHkbGWCV9PArLBw0rNqmdy56VOTRNTifBSpkC796Oez1lMFU+yDtkElbcrRSIlS5hRFqpsaFA==" diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 700707ce..a25f3010 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -1,7 +1,17 @@ # https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates version: 2 +enable-beta-ecosystems: true updates: - package-ecosystem: "github-actions" directory: "/" # Location of package manifests schedule: interval: "weekly" + + - package-ecosystem: "julia" + directories: + # Location of Julia projects + - "/" + - "/examples/GettingStarted" + - "/examples/SymbolicOptimalControl" + schedule: + interval: "weekly" diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 141d1572..ac3017a0 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -29,13 +29,16 @@ jobs: os: - ubuntu-latest group: - - "all" + - "layers" + - "vision" + - "piml" + - "core" uses: LuxDL/Lux.jl/.github/workflows/CommonCI.yml@main with: julia_version: ${{ matrix.version }} os: ${{ matrix.os }} project: "." - test_args: "BACKEND_GROUP=CPU BOLTZ_TEST_GROUP=${{ matrix.group }}" + test_args: "${{ matrix.group }}" downgrade: uses: LuxDL/Lux.jl/.github/workflows/CommonCI.yml@main @@ -44,4 +47,4 @@ jobs: os: ubuntu-latest project: "." downgrade_testing: true - test_args: "BACKEND_GROUP=CPU BOLTZ_TEST_GROUP=all" + test_args: "" diff --git a/.github/workflows/CompatHelper.yml b/.github/workflows/CompatHelper.yml deleted file mode 100644 index c4ce1826..00000000 --- a/.github/workflows/CompatHelper.yml +++ /dev/null @@ -1,44 +0,0 @@ -name: CompatHelper - -on: - schedule: - - cron: 0 0 * * * - workflow_dispatch: - -permissions: - contents: write - pull-requests: write - -jobs: - CompatHelper: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v6 - - uses: julia-actions/setup-julia@v2 - with: - version: 1 - - name: "Add the General registry via Git" - run: | - import Pkg - ENV["JULIA_PKG_SERVER"] = "" - Pkg.Registry.add("General") - shell: julia --color=yes {0} - - name: "Install CompatHelper" - run: | - import Pkg - name = "CompatHelper" - uuid = "aa819f21-2bde-4658-8897-bab36330d9b7" - version = "3" - Pkg.add(; name, uuid, version) - shell: julia --color=yes {0} - - name: "Run CompatHelper" - run: | - import CompatHelper - subdirs = ["", "docs", "test"] - append!(subdirs, joinpath.(("examples",), filter(p -> isdir(joinpath("examples", p)), readdir("examples")))) - CompatHelper.main(; subdirs) - shell: julia --color=yes {0} - working-directory: "./" - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - COMPATHELPER_PRIV: ${{ secrets.DOCUMENTER_KEY }} diff --git a/Project.toml b/Project.toml index 0712aecc..c4979bdc 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Boltz" uuid = "4544d5e4-abc5-4dea-817f-29e4c205d9c8" authors = ["Avik Pal and contributors"] -version = "1.7.1" +version = "1.8.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -56,7 +56,7 @@ BoltzTrackerExt = "Tracker" BoltzZygoteExt = "Zygote" [compat] -ADTypes = "1.10" +ADTypes = "1.15" ArgCheck = "2.3" Artifacts = "1.10, 1" ChainRulesCore = "1.25.1" @@ -73,7 +73,7 @@ LazyArtifacts = "1.10" Lux = "1.21.2" LuxCore = "1.2" LuxLib = "1.11" -MLDataDevices = "1.11.2" +MLDataDevices = "1.17.3" Markdown = "1.10" Metalhead = "0.9.5" NNlib = "0.9.30" @@ -84,7 +84,7 @@ ReactantCore = "0.1.15" Reexport = "1.2.2" ReverseDiff = "1.16.1" SafeTensors = "1.2" -Scratch = "1.2" +Scratch = "1.3" Setfield = "1.1.2" Static = "1.1.1" Statistics = "1.10" @@ -92,3 +92,6 @@ Tracker = "0.2.38" WeightInitializers = "1" Zygote = "0.7.7" julia = "1.10" + +[workspace] +projects = ["test", "docs"] diff --git a/examples/GettingStarted/Project.toml b/examples/GettingStarted/Project.toml index 57a90418..b601248d 100644 --- a/examples/GettingStarted/Project.toml +++ b/examples/GettingStarted/Project.toml @@ -5,6 +5,9 @@ Lux = "b2108857-7c20-44ae-9111-449ecde12c47" Metalhead = "dbeba491-748d-5e0e-a39e-b530a07fa0cc" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +[sources] +Boltz = {path = "../.."} + [compat] Boltz = "1" JLD2 = "0.5, 0.6" diff --git a/examples/SymbolicOptimalControl/Project.toml b/examples/SymbolicOptimalControl/Project.toml index 6d3c8b3b..84bbb13f 100644 --- a/examples/SymbolicOptimalControl/Project.toml +++ b/examples/SymbolicOptimalControl/Project.toml @@ -17,6 +17,9 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" SymbolicRegression = "8254be44-1295-4e6a-a16d-46603ac705cb" SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b" +[sources] +Boltz = {path = "../.."} + [compat] Boltz = "1" CairoMakie = "0.12, 0.13, 0.14, 0.15" @@ -24,12 +27,12 @@ ComponentArrays = "0.15.11" DynamicExpressions = "1.10, 2" Latexify = "0.16.2" Lux = "1.21.2" -MLJ = "0.20.3, 0.21" -Optimization = "4" +MLJ = "0.20.3, 0.21, 0.22" +Optimization = "4, 5" OptimizationOptimJL = "0.4" OptimizationOptimisers = "0.3.2" OrdinaryDiffEqVerner = "1" SciMLSensitivity = "7.57" Statistics = "1.10" SymbolicRegression = "1" -SymbolicUtils = "3" +SymbolicUtils = "3, 4" diff --git a/examples/SymbolicOptimalControl/main.jl b/examples/SymbolicOptimalControl/main.jl index fee7e50e..4d33c486 100644 --- a/examples/SymbolicOptimalControl/main.jl +++ b/examples/SymbolicOptimalControl/main.jl @@ -32,17 +32,9 @@ # ## Package Imports -using Lux, - Boltz, - ComponentArrays, - OrdinaryDiffEqVerner, - Optimization, - OptimizationOptimJL, - OptimizationOptimisers, - SciMLSensitivity, - Statistics, - Printf, - Random +using Lux, Boltz, ComponentArrays, Statistics, Printf, Random +using Optimization, + OptimizationOptimJL, OptimizationOptimisers, SciMLSensitivity, OrdinaryDiffEqVerner using DynamicExpressions, SymbolicRegression, MLJ, SymbolicUtils, Latexify using CairoMakie diff --git a/src/initialize.jl b/src/initialize.jl index e96817f8..b32db550 100644 --- a/src/initialize.jl +++ b/src/initialize.jl @@ -1,6 +1,5 @@ module InitializeModels -using ArgCheck: @argcheck using Functors: fmap using Random: Random, AbstractRNG diff --git a/src/piml/PIML.jl b/src/piml/PIML.jl index 31fc7fbb..0c6f3678 100644 --- a/src/piml/PIML.jl +++ b/src/piml/PIML.jl @@ -1,6 +1,5 @@ module PIML -using ArgCheck: @argcheck using Compat: @compat using ConcreteStructs: @concrete using Random: Random, AbstractRNG diff --git a/test/Project.toml b/test/Project.toml index a96fc7ec..4e815884 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,6 +1,6 @@ [deps] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" -CPUSummary = "2a0fbf3d-bb9c-48f3-b0a9-814d99fd7ab9" +Boltz = "4544d5e4-abc5-4dea-817f-29e4c205d9c8" ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" @@ -11,25 +11,25 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" -LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda" LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11" LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531" MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40" Metalhead = "dbeba491-748d-5e0e-a39e-b530a07fa0cc" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" +ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc" Pickle = "fbb45041-c46e-462f-888f-7c521cafbc2c" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" -ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" -Reexport = "189a3867-3050-52da-a836-e630ba90ab69" +Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[sources] +Boltz = {path = ".."} [compat] Aqua = "0.8.7" -CPUSummary = "0.2" ComponentArrays = "0.15.16" DataInterpolations = "8" Downloads = "1.6" @@ -40,18 +40,15 @@ ForwardDiff = "1" GPUArraysCore = "0.1.6, 0.2" JLD2 = "0.5, 0.6" Lux = "1.21.2" -LuxCUDA = "0.3.3" LuxLib = "1" LuxTestUtils = "2" MLDataDevices = "1.11.2" Metalhead = "0.9.5" NNlib = "0.9.30" +ParallelTestRunner = "2.1" Pickle = "0.3.5" Pkg = "1.10" Random = "1.10" -ReTestItems = "1.24.0" Reactant = "0.2.97" -Reexport = "1.2.2" StableRNGs = "1.0.2" Test = "1.10" -Zygote = "0.7.7" diff --git a/test/qa_tests.jl b/test/core/qa.jl similarity index 72% rename from test/qa_tests.jl rename to test/core/qa.jl index c67f0fd8..bbc3e651 100644 --- a/test/qa_tests.jl +++ b/test/core/qa.jl @@ -1,15 +1,11 @@ -@testitem "Aqua: Quality Assurance" tags = [:others] begin - using Aqua +using Lux, Aqua, ExplicitImports, Boltz, Test +@testset "Aqua: Quality Assurance" begin Aqua.test_all(Boltz; ambiguities=false) Aqua.test_ambiguities(Boltz; recursive=false) end -@testitem "Explicit Imports: Quality Assurance" tags = [:others] begin - using Lux: Lux - using Zygote: Zygote # Load all trigger packages - using ExplicitImports - +@testset "Explicit Imports: Quality Assurance" begin @test check_no_implicit_imports(Boltz; skip=(Base, Core, Lux)) === nothing @test check_no_stale_explicit_imports(Boltz) === nothing @test check_no_self_qualified_accesses(Boltz) === nothing diff --git a/test/layer_tests.jl b/test/layer_tests.jl deleted file mode 100644 index f70934bb..00000000 --- a/test/layer_tests.jl +++ /dev/null @@ -1,534 +0,0 @@ -# Only tests that are not run via `vision` or other higher-level test suites are -# included in this snippet. -@testitem "MLP" setup = [SharedTestSetup] tags = [:layers] begin - @testset "$(mode)" for (mode, aType, dev) in MODES - @testset "$(act)" for act in (tanh,) - @testset "$(nType)" for nType in (BatchNorm,) - norm = if nType === nothing - nType - elseif nType === BatchNorm - (i, ch, act; kwargs...) -> BatchNorm(ch, act; kwargs...) - elseif nType === GroupNorm - (i, ch, act; kwargs...) -> GroupNorm(ch, 2, act; kwargs...) - end - - model = Layers.MLP(2, (4, 4, 2), act; norm_layer=norm) - ps, st = Lux.setup(StableRNG(0), model) |> dev - st_test = Lux.testmode(st) - - x = randn(Float32, 2, 2) |> aType - - @test_gradients( - sumabs2first, - Constant(model), - x, - ps, - Constant(st); - atol=1e-3, - rtol=1e-3, - soft_fail=[AutoFiniteDiff()], - enzyme_set_runtime_activity=true - ) - - if test_reactant(mode) - set_reactant_backend!(mode) - rdev = reactant_device(; force=true) - - ps_ra, st_ra, x_ra = rdev((ps, st, x)) - st_ra_test = Lux.testmode(st_ra) - - @test @jit(model(x_ra, ps_ra, st_ra_test))[1] ≈ model(x, ps, st_test)[1] atol = - 1e-3 rtol = 1e-3 - - ∂x_ra, ∂ps_ra = - @jit(compute_reactant_gradient(model, x_ra, ps_ra, st_ra)) |> - cpu_device() - ∂x_zyg, ∂ps_zyg = - compute_zygote_gradient(model, x, ps, st) |> cpu_device() - @test check_approx(∂x_ra, ∂x_zyg; atol=1e-3, rtol=1e-3) - @test check_approx(∂ps_ra, ∂ps_zyg; atol=1e-3, rtol=1e-3) - end - end - end - end -end - -@testitem "Hamiltonian Neural Network" setup = [SharedTestSetup] tags = [:layers] begin - using ComponentArrays, ForwardDiff, Zygote, MLDataDevices, NNlib - - _remove_nothing(xs) = map(x -> x === nothing ? 0 : x, xs) - - @testset "$(mode): $(autodiff)" for (mode, aType, dev) in MODES, - autodiff in (nothing, AutoZygote(), AutoForwardDiff()) - - dev isa MLDataDevices.AbstractGPUDevice && - autodiff === AutoForwardDiff() && - continue - - hnn = Layers.HamiltonianNN{true}(Layers.MLP(2, (4, 4, 2), NNlib.gelu); autodiff) - ps, st = dev(Lux.setup(StableRNG(0), hnn)) - - x = aType(randn(Float32, 2, 4)) - - @test_throws ArgumentError hnn(x, ps, st) - - hnn = Layers.HamiltonianNN{true}(Layers.MLP(2, (4, 4, 1), NNlib.gelu); autodiff) - ps, st = dev(Lux.setup(StableRNG(0), hnn)) - ps_ca = dev(ComponentArray(cpu_device()(ps))) - - @test st.first_call - y, st = hnn(x, ps, st) - @test !st.first_call - - ∂x_zyg, ∂ps_zyg = Zygote.gradient( - (x, ps) -> sum(abs2, first(hnn(x, ps, st))), x, ps - ) - @test ∂x_zyg !== nothing - @test ∂ps_zyg !== nothing - if !(dev isa MLDataDevices.AbstractGPUDevice) - ∂ps_zyg = _remove_nothing(getdata(dev(ComponentArray(cpu_device()(∂ps_zyg))))) - ∂x_fd = ForwardDiff.gradient(x -> sum(abs2, first(hnn(x, ps, st))), x) - ∂ps_fd = getdata( - ForwardDiff.gradient(ps -> sum(abs2, first(hnn(x, ps, st))), ps_ca) - ) - - @test ∂x_zyg ≈ ∂x_fd atol = 1e-3 rtol = 1e-3 - @test ∂ps_zyg ≈ ∂ps_fd atol = 1e-3 rtol = 1e-3 - end - - st = Lux.initialstates(StableRNG(0), hnn) |> dev - st_test = Lux.testmode(st) - - @test st.first_call - y, st = hnn(x, ps_ca, st) - @test !st.first_call - - ∂x_zyg, ∂ps_zyg = Zygote.gradient( - (x, ps) -> sum(abs2, first(hnn(x, ps, st))), x, ps_ca - ) - @test ∂x_zyg !== nothing - @test ∂ps_zyg !== nothing - if !(dev isa MLDataDevices.AbstractGPUDevice) - ∂ps_zyg = _remove_nothing(getdata(dev(ComponentArray(cpu_device()(∂ps_zyg))))) - ∂x_fd = ForwardDiff.gradient(x -> sum(abs2, first(hnn(x, ps_ca, st))), x) - ∂ps_fd = getdata( - ForwardDiff.gradient(ps -> sum(abs2, first(hnn(x, ps, st))), ps_ca) - ) - - @test ∂x_zyg ≈ ∂x_fd atol = 1e-3 rtol = 1e-3 - @test ∂ps_zyg ≈ ∂ps_fd atol = 1e-3 rtol = 1e-3 - end - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra, x_ra = rdev((ps, st, x)) - st_ra_test = Lux.testmode(st_ra) - - @test @jit(hnn(x_ra, ps_ra, st_ra_test))[1] ≈ hnn(x, ps, st_test)[1] atol = 1e-3 rtol = - 1e-3 - - ∂x_ra, ∂ps_ra = - @jit(compute_reactant_gradient(hnn, x_ra, ps_ra, st_ra)) |> cpu_device() - ∂x_zyg, ∂ps_zyg = compute_zygote_gradient(hnn, x, ps, st) |> cpu_device() - - @test check_approx(∂x_ra, ∂x_zyg; atol=1e-3, rtol=1e-3) - @test check_approx(∂ps_ra, ∂ps_zyg; atol=1e-3, rtol=1e-3) - end - end -end - -@testitem "Tensor Product Layer" setup = [SharedTestSetup] tags = [:layers] begin - @testset "$(mode)" for (mode, aType, dev) in MODES - @testset "$(basis)" for basis in ( - Basis.Chebyshev, - Basis.Sin, - Basis.Cos, - Basis.Fourier, - Basis.Legendre, - Basis.Polynomial, - ) - tensor_project = Layers.TensorProductLayer([basis(n + 2) for n in 1:3], 4) - ps, st = dev(Lux.setup(StableRNG(0), tensor_project)) - - x = aType(tanh.(randn(Float32, 2, 4, 5))) - - @test_throws ArgumentError tensor_project(x, ps, st) - - x = aType(tanh.(randn(Float32, 2, 3, 5))) - - y, st = tensor_project(x, ps, st) - @test size(y) == (2, 4, 5) - - @test_gradients( - sumabs2first, - Constant(tensor_project), - x, - ps, - Constant(st); - atol=1e-3, - rtol=1e-3, - skip_backends=[AutoEnzyme()] - ) - - if test_reactant(mode) - set_reactant_backend!(mode) - - # XXX: Currently causes some issues with tracing - basis == Basis.Legendre && continue - - rdev = reactant_device(; force=true) - - x_ra = rdev(x) - ps_ra, st_ra = rdev((ps, st)) - st_ra_test = Lux.testmode(st_ra) - - @test @jit(tensor_project(x_ra, ps_ra, st_ra_test))[1] ≈ - tensor_project(x, ps, st)[1] atol = 1e-3 rtol = 1e-3 - - ∂x_ra, ∂ps_ra = - @jit(compute_reactant_gradient(tensor_project, x_ra, ps_ra, st_ra)) |> - cpu_device() - ∂x_zyg, ∂ps_zyg = - compute_zygote_gradient(tensor_project, x, ps, st) |> cpu_device() - - @test check_approx(∂x_ra, ∂x_zyg; atol=1e-3, rtol=1e-3) - @test check_approx(∂ps_ra, ∂ps_zyg; atol=1e-3, rtol=1e-3) - end - end - end -end - -@testitem "Basis Functions" setup = [SharedTestSetup] tags = [:layers] begin - @testset "$(mode)" for (mode, aType, dev) in MODES - @testset "$(basis)" for basis in ( - Basis.Chebyshev, - Basis.Sin, - Basis.Cos, - Basis.Fourier, - Basis.Legendre, - Basis.Polynomial, - ) - x = aType(tanh.(randn(Float32, 2, 4))) - grid = aType(collect(1:3)) - - fn1 = basis(3) - @test size(fn1(x)) == (3, 2, 4) - @test size(fn1(x, grid)) == (3, 2, 4) - - fn2 = basis(3; dim=2) - @test size(fn2(x)) == (2, 3, 4) - @test size(fn2(x, grid)) == (2, 3, 4) - - fn3 = basis(3; dim=3) - @test size(fn3(x)) == (2, 4, 3) - @test size(fn3(x, grid)) == (2, 4, 3) - - fn4 = basis(3; dim=4) - @test_throws ArgumentError fn4(x) - - grid2 = aType(1:5) - @test_throws ArgumentError fn4(x, grid2) - - if test_reactant(mode) - set_reactant_backend!(mode) - - # XXX: Currently causes some issues with tracing - basis == Basis.Legendre && continue - - rdev = reactant_device(; force=true) - - x_ra = rdev(x) - grid_ra = rdev(grid) - - @test @jit(fn1(x_ra)) ≈ fn1(x) atol = 1e-3 rtol = 1e-3 - @test @jit(fn1(x_ra, grid_ra)) ≈ fn1(x, grid) atol = 1e-3 rtol = 1e-3 - - @test @jit(fn2(x_ra)) ≈ fn2(x) atol = 1e-3 rtol = 1e-3 - @test @jit(fn2(x_ra, grid_ra)) ≈ fn2(x, grid) atol = 1e-3 rtol = 1e-3 - - @test @jit(fn3(x_ra)) ≈ fn3(x) atol = 1e-3 rtol = 1e-3 - @test @jit(fn3(x_ra, grid_ra)) ≈ fn3(x, grid) atol = 1e-3 rtol = 1e-3 - end - end - end -end - -@testitem "Spline Layer" setup = [SharedTestSetup] tags = [:integration] begin - using ComponentArrays, DataInterpolations, ForwardDiff, Zygote, MLDataDevices - - @testset "$(mode)" for (mode, aType, dev) in MODES - dev isa MLDataDevices.AbstractGPUDevice && continue - - @testset "$(spl): train_grid $(train_grid), dims $(dims)" for spl in ( - ConstantInterpolation, - LinearInterpolation, - QuadraticInterpolation, - # QuadraticSpline, # XXX: DataInterpolations.jl broke it again!!! - CubicSpline, - ), - train_grid in (true, false), - dims in ((), (8,)) - - spline = Layers.SplineLayer(dims, 0.0f0, 1.0f0, 0.1f0, spl; train_grid) - ps, st = dev(Lux.setup(StableRNG(0), spline)) - ps_ca = dev(ComponentArray(cpu_device()(ps))) - - x = aType(rand(Float32, 4)) - - y, st = spline(x, ps, st) - @test size(y) == (dims..., 4) - - y, st = spline(x, ps_ca, st) - @test size(y) == (dims..., 4) - - ∂x, ∂ps = Zygote.gradient((x, ps) -> sum(abs2, first(spline(x, ps, st))), x, ps) - spl !== ConstantInterpolation && @test ∂x !== nothing - @test ∂ps !== nothing - - ∂x_fd = ForwardDiff.gradient(x -> sum(abs2, first(spline(x, ps, st))), x) - ∂ps_fd = ForwardDiff.gradient(ps -> sum(abs2, first(spline(x, ps, st))), ps_ca) - - spl !== ConstantInterpolation && @test ∂x ≈ ∂x_fd atol = 1e-3 rtol = 1e-3 - - @test ∂ps.saved_points ≈ ∂ps_fd.saved_points atol = 1e-3 rtol = 1e-3 - if train_grid - if ∂ps.grid === nothing - @test_softfail all(Base.Fix1(isapprox, 0), ∂ps_fd.grid) - else - @test ∂ps.grid ≈ ∂ps_fd.grid atol = 1e-3 rtol = 1e-3 - end - end - end - end -end - -@testitem "Periodic Embedding" setup = [SharedTestSetup] tags = [:layers] begin - @testset "$(mode)" for (mode, aType, dev) in MODES - layer = Layers.PeriodicEmbedding([2, 3], [4.0, π / 5]) - ps, st = dev(Lux.setup(StableRNG(0), layer)) - x = aType(randn(StableRNG(0), 6, 4, 3, 2)) - Δx = aType([0.0, 12.0, -2π / 5, 0.0, 0.0, 0.0]) - - val = Array(layer(x, ps, st)[1]) - shifted_val = Array(layer(x .+ Δx, ps, st)[1]) - - @test all(val[1:4, :, :, :] .== shifted_val[1:4, :, :, :]) && all( - isapprox.(val[5:8, :, :, :], shifted_val[5:8, :, :, :]; atol=5 * eps(Float32)) - ) - - @test_gradients( - sumabs2first, - Constant(layer), - x, - ps, - Constant(st); - atol=1.0f-3, - rtol=1.0f-3, - enzyme_set_runtime_activity=true - ) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra, x_ra = rdev((ps, st, x)) - st_ra_test = Lux.testmode(st_ra) - - @test @jit(layer(x_ra, ps_ra, st_ra_test))[1] ≈ layer(x, ps, st)[1] atol = 1e-3 rtol = - 1e-3 - - ∂x_ra, ∂ps_ra = - @jit(compute_reactant_gradient(layer, x_ra, ps_ra, st_ra)) |> cpu_device() - ∂x_zyg, ∂ps_zyg = compute_zygote_gradient(layer, x, ps, st) |> cpu_device() - - @test check_approx(∂x_ra, ∂x_zyg; atol=1e-3, rtol=1e-3) - @test check_approx(∂ps_ra, ∂ps_zyg; atol=1e-3, rtol=1e-3) - end - end -end - -@testitem "Dynamic Expressions Layer" setup = [SharedTestSetup] tags = [:integration] begin - using DynamicExpressions, ForwardDiff, ComponentArrays - - operators = OperatorEnum(; binary_operators=[+, -, *], unary_operators=[cos]) - - x1 = Node(; feature=1) - x2 = Node(; feature=2) - - expr_1 = x1 * cos(x2 - 3.2) - expr_2 = x2 - x1 * x2 + 2.5 - 1.0 * x1 - - for exprs in ((expr_1,), (expr_1, expr_2), ([expr_1, expr_2],)) - layer = Layers.DynamicExpressionsLayer(operators, exprs...) - ps, st = Lux.setup(StableRNG(0), layer) - - x = [ - 1.0f0 2.0f0 3.0f0 - 4.0f0 5.0f0 6.0f0 - ] - - y, st_ = layer(x, ps, st) - @test eltype(y) == Float32 - - @test_gradients( - sumabs2first, - Constant(layer), - x, - ps, - Constant(st); - atol=1.0f-3, - rtol=1.0f-3, - skip_backends=[AutoEnzyme()] - ) - - # Particular ForwardDiff dispatches - ps_ca = ComponentArray(ps) - dps_ca = ForwardDiff.gradient(ps_ca) do ps_ - sum(abs2, first(layer(x, ps_, st))) - end - dx = ForwardDiff.gradient(x) do x_ - sum(abs2, first(layer(x_, ps, st))) - end - dxps = ForwardDiff.gradient(ComponentArray(; x, ps)) do ca - sum(abs2, first(layer(ca.x, ca.ps, st))) - end - - @test dx ≈ dxps.x atol = 1.0f-3 rtol = 1.0f-3 - @test dps_ca ≈ dxps.ps atol = 1.0f-3 rtol = 1.0f-3 - - x = Float64.(x) - y, st_ = layer(x, ps, st) - @test eltype(y) == Float64 - - @test_gradients( - sumabs2first, - Constant(layer), - x, - ps, - Constant(st); - atol=1.0e-3, - rtol=1.0e-3, - skip_backends=[AutoEnzyme()] - ) - end - - @testset "$(mode)" for (mode, aType, dev) in MODES - layer = Layers.DynamicExpressionsLayer(operators, expr_1) - ps, st = dev(Lux.setup(StableRNG(0), layer)) - - x = aType([ - 1.0f0 2.0f0 3.0f0 - 4.0f0 5.0f0 6.0f0 - ]) - - if dev isa MLDataDevices.AbstractGPUDevice - @test_throws ArgumentError layer(x, ps, st) - end - end -end - -@testitem "Positive Definite Container" setup = [SharedTestSetup] tags = [:layers] begin - @testset "$(mode)" for (mode, aType, dev) in MODES - model = Layers.MLP(2, (4, 4, 2), gelu) - pd = Layers.PositiveDefinite(model; in_dims=2) - ps, st = dev(Lux.setup(StableRNG(0), pd)) - - x = aType(randn(StableRNG(0), Float32, 2, 2)) - x0 = aType(zeros(Float32, 2)) - - y, _ = pd(x, ps, st) - z, _ = model(x, ps, st.model) - z0, _ = model(x0, ps, st.model) - y_by_hand = sum(abs2, z .- z0; dims=1) .+ sum(abs2, x .- x0; dims=1) - - @test maximum(abs, y - y_by_hand) < 1.0f-8 - - @test_gradients( - sumabs2first, - Constant(pd), - x, - ps, - Constant(st); - atol=1.0f-3, - rtol=1.0f-3, - broken_backends=[AutoEnzyme()] - ) - - pd2 = Layers.PositiveDefinite(model, ones(2)) - ps, st = dev(Lux.setup(StableRNG(0), pd2)) - - x0 = aType(ones(Float32, 2)) - y, _ = pd2(x0, ps, st) - - @test maximum(abs, y) < 1.0f-8 - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - pd = Layers.PositiveDefinite(model; in_dims=2) - ps, st = dev(Lux.setup(StableRNG(0), pd)) - x = aType(randn(StableRNG(0), Float32, 2, 2)) - ps_ra, st_ra, x_ra = rdev((ps, st, x)) - st_ra_test = Lux.testmode(st_ra) - - @test @jit(pd(x_ra, ps_ra, st_ra_test))[1] ≈ pd(x, ps, st)[1] atol = 1e-3 rtol = - 1e-3 - - ∂x_ra, ∂ps_ra = - @jit(compute_reactant_gradient(pd, x_ra, ps_ra, st_ra)) |> cpu_device() - ∂x_zyg, ∂ps_zyg = compute_zygote_gradient(pd, x, ps, st) |> cpu_device() - - @test check_approx(∂x_ra, ∂x_zyg; atol=1e-3, rtol=1e-3) - @test check_approx(∂ps_ra, ∂ps_zyg; atol=1e-3, rtol=1e-3) - end - end -end - -@testitem "ShiftTo Container" setup = [SharedTestSetup] tags = [:layers] begin - @testset "$(mode)" for (mode, aType, dev) in MODES - model = Layers.MLP(2, (4, 4, 2), gelu) - shiftto = Layers.ShiftTo(model, ones(Float32, 2), zeros(Float32, 2)) - ps, st = Lux.setup(StableRNG(0), shiftto) |> dev - - y0, _ = shiftto(st.in_val, ps, st) - @test maximum(abs, y0) < 1.0f-8 - - x = randn(StableRNG(0), Float32, 2, 2) |> aType - - @test_gradients( - sumabs2first, - Constant(shiftto), - x, - ps, - Constant(st); - atol=1.0f-3, - rtol=1.0f-3, - broken_backends=[AutoEnzyme()] - ) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra, x_ra = rdev((ps, st, x)) - st_ra_test = Lux.testmode(st_ra) - - @test @jit(shiftto(x_ra, ps_ra, st_ra_test))[1] ≈ shiftto(x, ps, st)[1] atol = - 1e-3 rtol = 1e-3 - - ∂x_ra, ∂ps_ra = - @jit(compute_reactant_gradient(shiftto, x_ra, ps_ra, st_ra)) |> cpu_device() - ∂x_zyg, ∂ps_zyg = compute_zygote_gradient(shiftto, x, ps, st) |> cpu_device() - - @test check_approx(∂x_ra, ∂x_zyg; atol=1e-3, rtol=1e-3) - @test check_approx(∂ps_ra, ∂ps_zyg; atol=1e-3, rtol=1e-3) - end - end -end diff --git a/test/layers/basis_functions.jl b/test/layers/basis_functions.jl new file mode 100644 index 00000000..d904397d --- /dev/null +++ b/test/layers/basis_functions.jl @@ -0,0 +1,52 @@ +using Reactant, Boltz, Lux, StableRNGs, Test + +include("../testutils.jl") + +dev = reactant_device(; force=true) + +# Test for each basis type +@testset "$(basis)" for basis in ( + Basis.Chebyshev, + Basis.Sin, + Basis.Cos, + Basis.Fourier, + Basis.Polynomial, + # Basis.Legendre # TODO: tracing +) + x = tanh.(randn(Float32, 2, 4)) + x_ra = x |> dev + grid = collect(Float32, 1:3) + grid_ra = grid |> dev + + # Test with default dimension (dim=1) + fn1 = basis(3) + @test size(fn1(x)) == (3, 2, 4) + @test size(fn1(x, grid)) == (3, 2, 4) + + # Test with dim=2 + fn2 = basis(3; dim=2) + @test size(fn2(x)) == (2, 3, 4) + @test size(fn2(x, grid)) == (2, 3, 4) + + # Test with dim=3 + fn3 = basis(3; dim=3) + @test size(fn3(x)) == (2, 4, 3) + @test size(fn3(x, grid)) == (2, 4, 3) + + # Test dimension error + fn4 = basis(3; dim=4) + @test_throws ArgumentError fn4(x) + + grid2 = collect(Float32, 1:5) + @test_throws ArgumentError fn4(x, grid2) + + # Test Reactant compilation + @test @jit(fn1(x_ra)) ≈ fn1(x) atol = 1e-3 rtol = 1e-3 + @test @jit(fn1(x_ra, grid_ra)) ≈ fn1(x, grid) atol = 1e-3 rtol = 1e-3 + + @test @jit(fn2(x_ra)) ≈ fn2(x) atol = 1e-3 rtol = 1e-3 + @test @jit(fn2(x_ra, grid_ra)) ≈ fn2(x, grid) atol = 1e-3 rtol = 1e-3 + + @test @jit(fn3(x_ra)) ≈ fn3(x) atol = 1e-3 rtol = 1e-3 + @test @jit(fn3(x_ra, grid_ra)) ≈ fn3(x, grid) atol = 1e-3 rtol = 1e-3 +end diff --git a/test/layers/dynamic_expressions.jl b/test/layers/dynamic_expressions.jl new file mode 100644 index 00000000..73658227 --- /dev/null +++ b/test/layers/dynamic_expressions.jl @@ -0,0 +1,68 @@ +using Boltz, Lux, StableRNGs, Test, LuxTestUtils +using DynamicExpressions, ForwardDiff, ComponentArrays + +include("../testutils.jl") + +operators = OperatorEnum(; binary_operators=[+, -, *], unary_operators=[cos]) + +x1 = Node(; feature=1) +x2 = Node(; feature=2) + +expr_1 = x1 * cos(x2 - 3.2) +expr_2 = x2 - x1 * x2 + 2.5 - 1.0 * x1 + +@testset "expressions: $(exprs)" for exprs in + ((expr_1,), (expr_1, expr_2), ([expr_1, expr_2],)) + layer = Layers.DynamicExpressionsLayer(operators, exprs...) + ps, st = Lux.setup(StableRNG(0), layer) + + x = Float32[ + 1.0 2.0 3.0 + 4.0 5.0 6.0 + ] + + y, st_ = layer(x, ps, st) + @test eltype(y) == Float32 + + @test_gradients( + TestUtils.sumabs2first, + Constant(layer), + x, + ps, + Constant(st); + atol=1.0f-2, + rtol=1.0f-2, + skip_backends=[AutoEnzyme()] + ) + + # Test ForwardDiff dispatches + ps_ca = ComponentArray(ps) + dps_ca = ForwardDiff.gradient(ps_ca) do ps_ + sum(abs2, first(layer(x, ps_, st))) + end + dx = ForwardDiff.gradient(x) do x_ + sum(abs2, first(layer(x_, ps, st))) + end + dxps = ForwardDiff.gradient(ComponentArray(; x, ps)) do ca + sum(abs2, first(layer(ca.x, ca.ps, st))) + end + + @test dx ≈ dxps.x atol = 1.0f-3 rtol = 1.0f-3 + @test dps_ca ≈ dxps.ps atol = 1.0f-3 rtol = 1.0f-3 + + # Test with Float64 + x64 = Float64.(x) + y64, st_ = layer(x64, ps, st) + @test eltype(y64) == Float64 + + @test_gradients( + TestUtils.sumabs2first, + Constant(layer), + x64, + ps, + Constant(st); + atol=1.0e-2, + rtol=1.0e-2, + skip_backends=[AutoEnzyme()] + ) +end diff --git a/test/layers/hamiltonian_nn.jl b/test/layers/hamiltonian_nn.jl new file mode 100644 index 00000000..f6f8788f --- /dev/null +++ b/test/layers/hamiltonian_nn.jl @@ -0,0 +1,37 @@ +using Reactant, Boltz, Lux, StableRNGs, Test, LuxTestUtils, NNlib +using ComponentArrays, ForwardDiff + +include("../testutils.jl") + +dev = reactant_device(; force=true) + +# Test that HNN with wrong output dimensions throws error +hnn_bad = Layers.HamiltonianNN{true}(Layers.MLP(2, (4, 4, 2), NNlib.gelu); autodiff=nothing) +ps, st = Lux.setup(StableRNG(0), hnn_bad) +x = randn(Float32, 2, 4) +@test_throws ArgumentError hnn_bad(x, ps, st) + +# Test HNN with correct output dimensions +hnn = Layers.HamiltonianNN{true}(Layers.MLP(2, (4, 4, 1), NNlib.gelu); autodiff=nothing) +ps, st = Lux.setup(StableRNG(0), hnn) +ps_ra, st_ra = (ps, st) |> dev + +x = randn(Float32, 2, 4) +x_ra = x |> dev + +@test st.first_call +y, st = hnn(x, ps, st) +@test !st.first_call + +st_test = Lux.testmode(st) +st_ra_test = st_test |> dev + +# Test Reactant forward pass +@test @jit(hnn(x_ra, ps_ra, st_ra_test))[1] ≈ hnn(x, ps, st_test)[1] atol = 1e-3 rtol = 1e-3 + +dx_ra, dps_ra = TestUtils.compute_reactant_gradient(hnn, x_ra, ps_ra, st_ra) +# TODO: Fix this?? Batching of autodiff calls +# dx_fd, dps_fd = TestUtils.compute_reactant_gradient_fd(hnn, x_ra, ps_ra, st_ra) + +# @test dx_ra ≈ dx_fd atol = 1e-3 rtol = 1e-3 +# @test LuxTestUtils.check_approx(dps_ra, dps_fd; atol=1e-3, rtol=1e-3) diff --git a/test/layers/mlp.jl b/test/layers/mlp.jl new file mode 100644 index 00000000..ac4aac2a --- /dev/null +++ b/test/layers/mlp.jl @@ -0,0 +1,36 @@ +using Reactant, Boltz, Lux, StableRNGs, Test, LuxTestUtils + +include("../testutils.jl") + +dev = reactant_device(; force=true) + +act = tanh +norm = (i, ch, act; kwargs...) -> BatchNorm(ch, act; kwargs...) +norm_2 = + (i, ch, act; kwargs...) -> + BatchNorm(ch, act; kwargs..., use_decomposed_implementation=true) + +model = Layers.MLP(2, (4, 4, 2), act; norm_layer=norm) +model_fd = Layers.MLP(2, (4, 4, 2), act; norm_layer=norm_2) + +ps, st = Lux.setup(StableRNG(0), model) +ps_ra, st_ra = (ps, st) |> dev + +x = randn(Float32, 2, 2) +x_ra = x |> dev + +st_test = Lux.testmode(st) +st_ra_test = st_test |> dev + +@test_gradients( + TestUtils.sumabs2first, Constant(model), x, ps, Constant(st); atol=1e-3, rtol=1e-3, +) + +@test @jit(model(x_ra, ps_ra, st_ra_test))[1] ≈ model(x, ps, st_test)[1] atol = 1e-3 rtol = + 1e-3 + +dx_ra, dps_ra = TestUtils.compute_reactant_gradient(model, x_ra, ps_ra, st_ra) +dx_fd, dps_fd = TestUtils.compute_reactant_gradient_fd(model_fd, x_ra, ps_ra, st_ra) + +@test dx_ra ≈ dx_fd atol = 1e-3 rtol = 1e-3 +@test LuxTestUtils.check_approx(dps_ra, dps_fd; atol=1e-3, rtol=1e-3) diff --git a/test/layers/periodic_embedding.jl b/test/layers/periodic_embedding.jl new file mode 100644 index 00000000..18bbff3a --- /dev/null +++ b/test/layers/periodic_embedding.jl @@ -0,0 +1,43 @@ +using Reactant, Boltz, Lux, StableRNGs, Test, LuxTestUtils + +include("../testutils.jl") + +dev = reactant_device(; force=true) + +layer = Layers.PeriodicEmbedding([2, 3], [4.0, π / 5]) +ps, st = Lux.setup(StableRNG(0), layer) +ps_ra, st_ra = (ps, st) |> dev + +x = randn(StableRNG(0), Float32, 6, 4, 3, 2) +x_ra = x |> dev +Δx = Float32[0.0, 12.0, -2π / 5, 0.0, 0.0, 0.0] + +# Test periodicity +val = Array(layer(x, ps, st)[1]) +shifted_val = Array(layer(x .+ Δx, ps, st)[1]) + +@test all(val[1:4, :, :, :] .== shifted_val[1:4, :, :, :]) +@test all(isapprox.(val[5:8, :, :, :], shifted_val[5:8, :, :, :]; atol=1e-5)) + +@test_gradients( + TestUtils.sumabs2first, + Constant(layer), + x, + ps, + Constant(st); + atol=1.0f-3, + rtol=1.0f-3, + broken_backends=[AutoEnzyme()] +) + +st_test = Lux.testmode(st) +st_ra_test = st_test |> dev + +@test @jit(layer(x_ra, ps_ra, st_ra_test))[1] ≈ layer(x, ps, st_test)[1] atol = 1e-3 rtol = + 1e-3 + +dx_ra, dps_ra = TestUtils.compute_reactant_gradient(layer, x_ra, ps_ra, st_ra) +dx_fd, dps_fd = TestUtils.compute_reactant_gradient_fd(layer, x_ra, ps_ra, st_ra) + +@test dx_ra ≈ dx_fd atol = 1e-3 rtol = 1e-3 +@test LuxTestUtils.check_approx(dps_ra, dps_fd; atol=1e-3, rtol=1e-3) diff --git a/test/layers/positive_definite.jl b/test/layers/positive_definite.jl new file mode 100644 index 00000000..6d3b4f5b --- /dev/null +++ b/test/layers/positive_definite.jl @@ -0,0 +1,54 @@ +using Reactant, Boltz, Lux, StableRNGs, Test, LuxTestUtils + +include("../testutils.jl") + +dev = reactant_device(; force=true) + +model = Layers.MLP(2, (4, 4, 2), gelu) +pd = Layers.PositiveDefinite(model; in_dims=2) +ps, st = Lux.setup(StableRNG(0), pd) +ps_ra, st_ra = (ps, st) |> dev + +x = randn(StableRNG(0), Float32, 2, 2) +x_ra = x |> dev +x0 = zeros(Float32, 2) + +# Verify positive definite property by hand +y, _ = pd(x, ps, st) +z, _ = model(x, ps, st.model) +z0, _ = model(x0, ps, st.model) +y_by_hand = sum(abs2, z .- z0; dims=1) .+ sum(abs2, x .- x0; dims=1) + +@test maximum(abs, y - y_by_hand) < 1.0f-8 + +@test_gradients( + TestUtils.sumabs2first, + Constant(pd), + x, + ps, + Constant(st); + atol=1.0f-3, + rtol=1.0f-3, + broken_backends=[AutoEnzyme()] +) + +# Test with explicit reference point +pd2 = Layers.PositiveDefinite(model, ones(2)) +ps2, st2 = Lux.setup(StableRNG(0), pd2) + +x0_ones = ones(Float32, 2) +y2, _ = pd2(x0_ones, ps2, st2) + +@test maximum(abs, y2) < 1.0f-8 + +# Test Reactant +st_test = Lux.testmode(st) +st_ra_test = st_test |> dev + +@test @jit(pd(x_ra, ps_ra, st_ra_test))[1] ≈ pd(x, ps, st_test)[1] atol = 1e-3 rtol = 1e-3 + +dx_ra, dps_ra = TestUtils.compute_reactant_gradient(pd, x_ra, ps_ra, st_ra) +dx_fd, dps_fd = TestUtils.compute_reactant_gradient_fd(pd, x_ra, ps_ra, st_ra) + +@test dx_ra ≈ dx_fd atol = 1e-3 rtol = 1e-3 +@test LuxTestUtils.check_approx(dps_ra, dps_fd; atol=1e-3, rtol=1e-3) diff --git a/test/layers/shiftto.jl b/test/layers/shiftto.jl new file mode 100644 index 00000000..db4cd801 --- /dev/null +++ b/test/layers/shiftto.jl @@ -0,0 +1,40 @@ +using Reactant, Boltz, Lux, StableRNGs, Test, LuxTestUtils + +include("../testutils.jl") + +dev = reactant_device(; force=true) + +model = Layers.MLP(2, (4, 4, 2), gelu) +shiftto = Layers.ShiftTo(model, ones(Float32, 2), zeros(Float32, 2)) + +ps, st = Lux.setup(StableRNG(0), shiftto) +ps_ra, st_ra = (ps, st) |> dev + +y0, _ = shiftto(st.in_val, ps, st) +@test maximum(abs, y0) < 1.0f-8 + +x = randn(StableRNG(0), Float32, 2, 2) +x_ra = x |> dev + +@test_gradients( + TestUtils.sumabs2first, + Constant(shiftto), + x, + ps, + Constant(st); + atol=1.0f-3, + rtol=1.0f-3, + broken_backends=[AutoEnzyme()] +) + +st_test = Lux.testmode(st) +st_ra_test = st_test |> dev + +@test @jit(shiftto(x_ra, ps_ra, st_ra_test))[1] ≈ shiftto(x, ps, st_test)[1] atol = 1e-3 rtol = + 1e-3 + +dx_ra, dps_ra = TestUtils.compute_reactant_gradient(shiftto, x_ra, ps_ra, st_ra) +dx_fd, dps_fd = TestUtils.compute_reactant_gradient_fd(shiftto, x_ra, ps_ra, st_ra) + +@test dx_ra ≈ dx_fd atol = 1e-3 rtol = 1e-3 +@test LuxTestUtils.check_approx(dps_ra, dps_fd; atol=1e-3, rtol=1e-3) diff --git a/test/layers/spline.jl b/test/layers/spline.jl new file mode 100644 index 00000000..f1945344 --- /dev/null +++ b/test/layers/spline.jl @@ -0,0 +1,31 @@ +using Boltz, Lux, StableRNGs, Test +using ComponentArrays, DataInterpolations, ForwardDiff + +# NOTE: Spline layer tests are CPU-only (no GPU/Reactant support) + +@testset "$(spl): train_grid=$(train_grid), dims=$(dims)" for spl in ( + ConstantInterpolation, LinearInterpolation, QuadraticInterpolation, CubicSpline + ), + train_grid in (true, false), + dims in ((), (8,)) + + spline = Layers.SplineLayer(dims, 0.0f0, 1.0f0, 0.1f0, spl; train_grid) + ps, st = Lux.setup(StableRNG(0), spline) + + x = rand(Float32, 4) + + y, st = spline(x, ps, st) + @test size(y) == (dims..., 4) + + # Test with ComponentArray + ps_ca = ComponentArray(ps) + + y, st = spline(x, ps_ca, st) + @test size(y) == (dims..., 4) + + ∂x_fd = ForwardDiff.gradient(x -> sum(abs2, first(spline(x, ps, st))), x) + ∂ps_fd = ForwardDiff.gradient(ps -> sum(abs2, first(spline(x, ps, st))), ps_ca) + + @test !all(iszero, ∂x_fd) skip = spl === ConstantInterpolation + @test !all(iszero, ∂ps_fd) +end diff --git a/test/layers/tensor_product.jl b/test/layers/tensor_product.jl new file mode 100644 index 00000000..ff1ce0c5 --- /dev/null +++ b/test/layers/tensor_product.jl @@ -0,0 +1,54 @@ +using Reactant, Boltz, Lux, StableRNGs, Test, LuxTestUtils + +include("../testutils.jl") + +dev = reactant_device(; force=true) + +# Test for each basis type +@testset "$(basis)" for basis in ( + Basis.Chebyshev, + Basis.Sin, + Basis.Cos, + Basis.Fourier, + # Basis.Legendre, # TODO: tracing + Basis.Polynomial, +) + tensor_project = Layers.TensorProductLayer([basis(n + 2) for n in 1:3], 4) + ps, st = Lux.setup(StableRNG(0), tensor_project) + ps_ra, st_ra = (ps, st) |> dev + + # Test dimension mismatch error + x_bad = tanh.(randn(Float32, 2, 4, 5)) + @test_throws ArgumentError tensor_project(x_bad, ps, st) + + x = tanh.(randn(Float32, 2, 3, 5)) + x_ra = x |> dev + + y, st = tensor_project(x, ps, st) + @test size(y) == (2, 4, 5) + + @test_gradients( + TestUtils.sumabs2first, + Constant(tensor_project), + x, + ps, + Constant(st); + atol=1e-3, + rtol=1e-3, + skip_backends=[AutoEnzyme()] + ) + + st_test = Lux.testmode(st) + st_ra_test = st_test |> dev + + @test @jit(tensor_project(x_ra, ps_ra, st_ra_test))[1] ≈ + tensor_project(x, ps, st_test)[1] atol = 1e-3 rtol = 1e-3 + + dx_ra, dps_ra = TestUtils.compute_reactant_gradient(tensor_project, x_ra, ps_ra, st_ra) + dx_fd, dps_fd = TestUtils.compute_reactant_gradient_fd( + tensor_project, x_ra, ps_ra, st_ra + ) + + @test dx_ra ≈ dx_fd atol = 1e-3 rtol = 1e-3 + @test LuxTestUtils.check_approx(dps_ra, dps_fd; atol=1e-3, rtol=1e-3) +end diff --git a/test/piml/transolver.jl b/test/piml/transolver.jl new file mode 100644 index 00000000..275a7b18 --- /dev/null +++ b/test/piml/transolver.jl @@ -0,0 +1,51 @@ +using Reactant, Boltz, Lux, StableRNGs, Test, LuxTestUtils + +include("../testutils.jl") + +rdev = reactant_device(; force=true) + +@testset "Base Model: $(use_rms_norm)" for use_rms_norm in (true, false) + model = PIML.Transolver(; + func_dim=6, spatial_dim=3, nheads=8, num_layers=1, out_dim=2, use_rms_norm + ) + ps, st = Lux.setup(StableRNG(0), model) + x = randn(Float32, 3, 32, 4) + fx = randn(Float32, 6, 32, 4) + + y, _ = model((x, fx), ps, st) + @test size(y) == (2, 32, 4) + + ps_ra, st_ra = rdev((ps, st)) + x_ra, fx_ra = rdev((x, fx)) + + @test @jit(model((x_ra, fx_ra), ps_ra, st_ra))[1] ≈ model((x, fx), ps, st)[1] atol = + 1e-3 rtol = 1e-3 +end + +@testset "Conditioned Model: $(use_rms_norm)" for use_rms_norm in (true, false) + hidden_dim = 128 + activation = gelu + preprocess = Parallel( + .+, + Dense(3 => hidden_dim, activation), + Chain( + Dense(6 => hidden_dim, activation), + WrappedFunction(x -> reshape(x, hidden_dim, 1, size(x, ndims(x)))), + ), + ) + + model = PIML.Transolver(; nheads=8, num_layers=1, out_dim=2, preprocess, use_rms_norm) + ps, st = Lux.setup(StableRNG(0), model) + + x = randn(Float32, 3, 32, 4) + fx = randn(Float32, 6, 4) # Global information about the mesh + + y, _ = model((x, fx), ps, st) + @test size(y) == (2, 32, 4) + + ps_ra, st_ra = rdev((ps, st)) + x_ra, fx_ra = rdev((x, fx)) + + @test @jit(model((x_ra, fx_ra), ps_ra, st_ra))[1] ≈ model((x, fx), ps, st)[1] atol = + 1e-3 rtol = 1e-3 +end diff --git a/test/piml_tests.jl b/test/piml_tests.jl deleted file mode 100644 index 7f435419..00000000 --- a/test/piml_tests.jl +++ /dev/null @@ -1,73 +0,0 @@ -@testitem "Transolver" setup = [SharedTestSetup] tags = [:piml] begin - for (mode, aType, dev) in MODES - @testset "Base Model: $(use_rms_norm)" for use_rms_norm in (true, false) - model = PIML.Transolver(; - func_dim=6, spatial_dim=3, nheads=8, num_layers=1, out_dim=2, use_rms_norm - ) - ps, st = dev(Lux.setup(Random.default_rng(), model)) - x = aType(randn(Float32, 3, 32, 4)) - fx = aType(randn(Float32, 6, 32, 4)) - - y, _ = model((x, fx), ps, st) - @test size(y) == (2, 32, 4) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra = rdev(cpu_device()((ps, st))) - x_ra, fx_ra = rdev(cpu_device()((x, fx))) - - Reactant.with_config(; - dot_general_precision=PrecisionConfig.HIGH, - convolution_precision=PrecisionConfig.HIGH, - ) do - @test @jit(model((x_ra, fx_ra), ps_ra, st_ra))[1] ≈ - model((x, fx), ps, st)[1] atol = 1e-3 rtol = 1e-3 - end - end - end - - @testset "Conditioned Model: $(use_rms_norm)" for use_rms_norm in (true, false) - hidden_dim = 128 - activation = gelu - preprocess = Parallel( - .+, - Dense(3 => hidden_dim, activation), - Chain( - Dense(6 => hidden_dim, activation), - WrappedFunction(x -> reshape(x, hidden_dim, 1, size(x, ndims(x)))), - ), - ) - - model = PIML.Transolver(; - nheads=8, num_layers=1, out_dim=2, preprocess, use_rms_norm - ) - ps, st = dev(Lux.setup(Random.default_rng(), model)) - - x = aType(randn(Float32, 3, 32, 4)) - fx = aType(randn(Float32, 6, 4)) # Global information about the mesh - - y, _ = model((x, fx), ps, st) - @test size(y) == (2, 32, 4) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra = rdev(cpu_device()((ps, st))) - x_ra, fx_ra = rdev(cpu_device()((x, fx))) - - Reactant.with_config(; - dot_general_precision=PrecisionConfig.HIGH, - convolution_precision=PrecisionConfig.HIGH, - ) do - @test @jit(model((x_ra, fx_ra), ps_ra, st_ra))[1] ≈ - model((x, fx), ps, st)[1] atol = 1e-3 rtol = 1e-3 - end - end - end - end -end diff --git a/test/runtests.jl b/test/runtests.jl index 7a961f7e..51893c86 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,85 +1,24 @@ -using ReTestItems, Pkg, CPUSummary, Test +using Boltz, ParallelTestRunner, Setfield -const ALL_BOTLZ_TEST_GROUPS = [ - "layers", "others", "vision", "vision_metalhead", "integration", "piml" -] +parsed_args = parse_args(@isdefined(TEST_ARGS) ? TEST_ARGS : ARGS) -function parse_test_args() - test_args_from_env = @isdefined(TEST_ARGS) ? TEST_ARGS : ARGS - test_args = Dict{String,String}() - for arg in test_args_from_env - if contains(arg, "=") - key, value = split(arg, "="; limit=2) - test_args[key] = value - end - end - @info "Parsed test args" test_args - return test_args -end - -const PARSED_TEST_ARGS = parse_test_args() - -const BOLTZ_TEST_REACTANT = parse( - Bool, lowercase(get(PARSED_TEST_ARGS, "BOLTZ_TEST_REACTANT", "true")) -) - -INPUT_TEST_GROUP = lowercase(get(PARSED_TEST_ARGS, "BOLTZ_TEST_GROUP", "all")) -const BOLTZ_TEST_GROUP = if startswith("!", INPUT_TEST_GROUP[1]) - exclude_group = lowercase.(split(INPUT_TEST_GROUP[2:end], ",")) - filter(x -> x ∉ exclude_group, ALL_BOTLZ_TEST_GROUPS) -else - [INPUT_TEST_GROUP] -end +testsuite = find_tests(@__DIR__) +delete!(testsuite, "testutils") +delete!(testsuite, "vision/testutils") -const BACKEND_GROUP = lowercase(get(PARSED_TEST_ARGS, "BACKEND_GROUP", "all")) -const EXTRA_PKGS = String[] - -(BACKEND_GROUP == "all" || BACKEND_GROUP == "amdgpu") && - !BOLTZ_TEST_REACTANT && - push!(EXTRA_PKGS, "AMDGPU") - -if !isempty(EXTRA_PKGS) - @info "Installing Extra Packages for testing" EXTRA_PKGS = EXTRA_PKGS - Pkg.add(EXTRA_PKGS) - Pkg.update() - Base.retry_load_extensions() - Pkg.instantiate() -end - -using Boltz - -const RETESTITEMS_NWORKERS = parse( - Int, - get( - ENV, - "RETESTITEMS_NWORKERS", - string(min(Int(CPUSummary.num_cores()), Sys.isapple() ? 2 : 4)), - ), +# Limit total jobs to 4 to avoid OOM on GPU +total_jobs = min( + something(parsed_args.jobs, ParallelTestRunner.default_njobs()), + length(keys(testsuite)), + 4, ) -const RETESTITEMS_NWORKER_THREADS = parse( - Int, - get( - ENV, - "RETESTITEMS_NWORKER_THREADS", - string(max(Int(CPUSummary.sys_threads()) ÷ RETESTITEMS_NWORKERS, 1)), - ), -) +@set! parsed_args.jobs = Some(total_jobs) -@testset "Boltz.jl Tests" begin - @testset for (i, tag) in enumerate(BOLTZ_TEST_GROUP) - withenv( - "BOLTZ_TEST_REACTANT" => BOLTZ_TEST_REACTANT, - "BACKEND_GROUP" => BACKEND_GROUP, - "XLA_REACTANT_GPU_MEM_FRACTION" => 1 / (RETESTITEMS_NWORKERS + 0.1), - ) do - ReTestItems.runtests( - Boltz; - tags=(tag == "all" ? nothing : [Symbol(tag)]), - testitem_timeout=2400, - nworkers=RETESTITEMS_NWORKERS, - nworker_threads=RETESTITEMS_NWORKER_THREADS, - ) - end - end +withenv( + "XLA_REACTANT_GPU_MEM_FRACTION" => 1 / (total_jobs + 0.1), + "XLA_REACTANT_GPU_PREALLOCATE" => false, + "JULIA_CUDA_HARD_MEMORY_LIMIT" => "$(100 / (total_jobs + 0.1))%", +) do + runtests(Boltz, parsed_args; testsuite) end diff --git a/test/shared_testsetup.jl b/test/shared_testsetup.jl deleted file mode 100644 index f23faee5..00000000 --- a/test/shared_testsetup.jl +++ /dev/null @@ -1,103 +0,0 @@ -@testsetup module SharedTestSetup - -import Reexport: @reexport -@reexport using Boltz, Lux, GPUArraysCore, LuxLib, LuxTestUtils, Random, StableRNGs, NNlib -using MLDataDevices, JLD2, Enzyme, Zygote -using LuxTestUtils: Constant, check_approx - -const BOLTZ_TEST_REACTANT = parse(Bool, lowercase(get(ENV, "BOLTZ_TEST_REACTANT", "true"))) - -const BACKEND_GROUP = lowercase(get(ENV, "BACKEND_GROUP", "all")) - -GPUArraysCore.allowscalar(false) - -if (BACKEND_GROUP == "all" || BACKEND_GROUP == "cuda") && !BOLTZ_TEST_REACTANT - using LuxCUDA -end - -if (BACKEND_GROUP == "all" || BACKEND_GROUP == "amdgpu") && !BOLTZ_TEST_REACTANT - using AMDGPU -end - -if BOLTZ_TEST_REACTANT - @reexport using Reactant -else - macro jit(ex) - quote - $(ex) - end - end - macro compile(ex) - quote - $(ex) - end - end - export @jit, @compile -end - -cpu_testing() = BACKEND_GROUP == "all" || BACKEND_GROUP == "cpu" -# Reactant will not work nicely if we load AMDGPU / CUDA -function cuda_testing() - return (BACKEND_GROUP == "all" || BACKEND_GROUP == "cuda") && - MLDataDevices.functional(CUDADevice) -end -function amdgpu_testing() - return (BACKEND_GROUP == "all" || BACKEND_GROUP == "amdgpu") && - MLDataDevices.functional(AMDGPUDevice) -end - -const MODES = begin - modes = [] - cpu_testing() && push!(modes, ("cpu", Array, CPUDevice())) - if !BOLTZ_TEST_REACTANT - cuda_testing() && push!(modes, ("cuda", CuArray, CUDADevice())) - amdgpu_testing() && push!(modes, ("amdgpu", ROCArray, AMDGPUDevice())) - else - if BACKEND_GROUP == "cuda" || BACKEND_GROUP == "all" - push!(modes, ("cuda", Array, CPUDevice())) - end - end - modes -end - -test_reactant(mode::String) = mode != "amdgpu" && BOLTZ_TEST_REACTANT -function set_reactant_backend!(mode::String) - if mode == "cuda" - Reactant.set_default_backend("gpu") - elseif mode == "cpu" - Reactant.set_default_backend("cpu") - else - error("Unknown mode $(mode)") - end -end - -sumabs2first(model, x, ps, st) = sum(abs2, first(model(x, ps, st))) - -function compute_reactant_gradient(model, x, ps, st) - return compute_reactant_gradient(sumabs2first, model, x, ps, st) -end - -function compute_reactant_gradient(f::F, model, x, ps, st) where {F} - res = Enzyme.gradient(Reverse, f, Const(model), x, ps, Const(st)) - return res[2], res[3] -end - -function compute_zygote_gradient(model, x, ps, st) - return compute_zygote_gradient(sumabs2first, model, x, ps, st) -end - -function compute_zygote_gradient(f::F, model, x, ps, st) where {F} - return Zygote.gradient((x, ps) -> f(model, x, ps, st), x, ps) -end - -export MODES, - BACKEND_GROUP, - test_reactant, - set_reactant_backend!, - compute_reactant_gradient, - compute_zygote_gradient, - check_approx, - Constant, - sumabs2first - -end diff --git a/test/testutils.jl b/test/testutils.jl new file mode 100644 index 00000000..0f87790b --- /dev/null +++ b/test/testutils.jl @@ -0,0 +1,29 @@ +module TestUtils + +using Reactant, Enzyme, Lux + +sumabs2first(model, x, ps, st) = sum(abs2, first(model(x, ps, st))) + +function compute_enzyme_gradient(model, x, ps, st) + return compute_enzyme_gradient(sumabs2first, model, x, ps, st) +end + +function compute_enzyme_gradient(f::F, model, x, ps, st) where {F} + res = Enzyme.gradient(Reverse, f, Const(model), x, ps, Const(st)) + return res[2], res[3] +end + +compute_reactant_gradient(args...) = @jit compute_enzyme_gradient(args...) + +function compute_reactant_gradient_fd(model, x, ps, st) + return compute_reactant_gradient_fd(sumabs2first, model, x, ps, st) +end + +function compute_reactant_gradient_fd(f::F, model, x, ps, st) where {F} + _, dx, dps, _ = @jit Reactant.TestUtils.finite_difference_gradient( + f, Const(model), f64(x), f64(ps), Const(st) + ) + return dx, dps +end + +end diff --git a/test/vision/alexnet.jl b/test/vision/alexnet.jl new file mode 100644 index 00000000..e7fffc76 --- /dev/null +++ b/test/vision/alexnet.jl @@ -0,0 +1,8 @@ +using Boltz, Test + +include("testutils.jl") + +@testset "AlexNet: pretrained: $(pretrained)" for pretrained in [true, false] + model = Vision.AlexNet(; pretrained) + VisionTestUtils.test_model(model; pretrained) +end diff --git a/test/vision/convmixer.jl b/test/vision/convmixer.jl new file mode 100644 index 00000000..7e591d88 --- /dev/null +++ b/test/vision/convmixer.jl @@ -0,0 +1,11 @@ +using Boltz, Test + +include("testutils.jl") + +@testset "ConvMixer: $(name): pretrained: $(pretrained)" for name in + [:small, :base, :large], + pretrained in [false] + + model = Vision.ConvMixer(name; pretrained) + VisionTestUtils.test_model(model; pretrained) +end diff --git a/test/vision/efficientnet.jl b/test/vision/efficientnet.jl new file mode 100644 index 00000000..3ddc8c27 --- /dev/null +++ b/test/vision/efficientnet.jl @@ -0,0 +1,16 @@ +using Boltz, Test + +include("testutils.jl") + +variants = if parse(Bool, get(ENV, "CI", "false")) + [:b0, :b1] +else + [:b0, :b1, :b2, :b3, :b4, :b5, :b6, :b7] +end + +@testset "EfficientNet: $(variant): pretrained: $(pretrained)" for variant in variants, + pretrained in [false, true] + + model = Vision.EfficientNet(variant; pretrained) + VisionTestUtils.test_model(model; pretrained) +end diff --git a/test/vision/googlenet.jl b/test/vision/googlenet.jl new file mode 100644 index 00000000..8782b712 --- /dev/null +++ b/test/vision/googlenet.jl @@ -0,0 +1,8 @@ +using Boltz, Test, Metalhead + +include("testutils.jl") + +@testset "GoogleNet: pretrained: $(pretrained)" for pretrained in [false] + model = Vision.GoogLeNet(; pretrained) + VisionTestUtils.test_model(model; pretrained) +end diff --git a/test/vision/mobilenet.jl b/test/vision/mobilenet.jl new file mode 100644 index 00000000..fb5a97e2 --- /dev/null +++ b/test/vision/mobilenet.jl @@ -0,0 +1,12 @@ +using Boltz, Test, Metalhead + +include("testutils.jl") + +@testset "MobileNet: $(name): pretrained: $(pretrained)" for name in [ + :v1, :v2, :v3_small, :v3_large + ], + pretrained in [false] + + model = Vision.MobileNet(name; pretrained) + VisionTestUtils.test_model(model; pretrained) +end diff --git a/test/vision/resnet.jl b/test/vision/resnet.jl new file mode 100644 index 00000000..f356ad67 --- /dev/null +++ b/test/vision/resnet.jl @@ -0,0 +1,12 @@ +using Boltz, Test, Metalhead + +include("testutils.jl") + +@testset "ResNet: $(depth): pretrained: $(pretrained)" for depth in [18, 34, 50, 101, 152], + pretrained in [false, true] + + pretrained && pkgversion(Metalhead) > v"0.9.4" && continue + + model = Vision.ResNet(depth; pretrained) + VisionTestUtils.test_model(model; pretrained) +end diff --git a/test/vision/resnext.jl b/test/vision/resnext.jl new file mode 100644 index 00000000..2fed1707 --- /dev/null +++ b/test/vision/resnext.jl @@ -0,0 +1,17 @@ +using Boltz, Test, Metalhead + +include("testutils.jl") + +@testset "ResNeXt: $(depth) : $(cardinality) : $(width) : pretrained: $(pretrained)" for ( + depth, cardinality, width + ) in [ + (50, 32, 4), (101, 32, 4), (152, 32, 4) + ], + pretrained in [false, true] + + depth == 152 && pretrained && continue + pretrained && pkgversion(Metalhead) > v"0.9.4" && continue + + model = Vision.ResNeXt(depth; pretrained, cardinality, base_width=width) + VisionTestUtils.test_model(model; pretrained) +end diff --git a/test/vision/squeezenet.jl b/test/vision/squeezenet.jl new file mode 100644 index 00000000..aa98e9cf --- /dev/null +++ b/test/vision/squeezenet.jl @@ -0,0 +1,8 @@ +using Boltz, Test, Metalhead + +include("testutils.jl") + +@testset "SqueezeNet: pretrained: $(pretrained)" for pretrained in [false, true] + model = Vision.SqueezeNet(; pretrained) + VisionTestUtils.test_model(model; pretrained) +end diff --git a/test/vision/testutils.jl b/test/vision/testutils.jl new file mode 100644 index 00000000..2b667221 --- /dev/null +++ b/test/vision/testutils.jl @@ -0,0 +1,69 @@ +module VisionTestUtils + +using Lux, Downloads, JLD2, Pickle, Reactant, StableRNGs, Reactant, Test + +function normalize_imagenet(data) + cmean = reshape(Float32[0.485, 0.456, 0.406], (1, 1, 3, 1)) + cstd = reshape(Float32[0.229, 0.224, 0.225], (1, 1, 3, 1)) + return (data .- cmean) ./ cstd +end + +# The images are normalized and saved +@load joinpath(@__DIR__, "../", "testimages", "monarch_color.jld2") monarch_color_224 monarch_color_256 +const MONARCH_224 = monarch_color_224 +const MONARCH_256 = monarch_color_256 + +const TEST_LBLS = readlines( + Downloads.download( + "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt" + ), +) + +function get_test_image(size) + size == 224 && return MONARCH_224 + size == 256 && return MONARCH_256 + return error("size must be 224 or 256") +end + +function imagenet_acctest(model, ps, st, dev; size=224) + ps = dev(ps) + st = dev(Lux.testmode(st)) + TEST_X = get_test_image(size) + x = dev(TEST_X) + + if dev isa ReactantDevice + res = @jit model(x, ps, st) + else + res = model(x, ps, st) + end + + ypred = vec(collect(cpu_device()(first(res)))) + top5 = TEST_LBLS[partialsortperm(ypred, 1:5; rev=true)] + return "monarch" in top5 +end + +function test_model(model; size=224, pretrained::Bool=false, nclasses::Int=1000) + ps, st = Lux.setup(StableRNGs.StableRNG(1234), model) + st = Lux.testmode(st) + img = get_test_image(size) + + res = first(model(img, ps, st)) + @test Base.size(res) == (nclasses, 1) + + if pretrained + @test imagenet_acctest(model, ps, st, CPUDevice(); size) + end + + GC.gc(true) + + rdev = reactant_device() + ps_ra, st_ra, img_ra = rdev((ps, st, img)) + res_ra = first(@jit model(img_ra, ps_ra, st_ra)) + @test res_ra ≈ res atol = 1e-3 rtol = 1e-3 + + if pretrained + @test imagenet_acctest(model, ps, st, rdev; size) + end +end + +end diff --git a/test/vision/vgg.jl b/test/vision/vgg.jl new file mode 100644 index 00000000..1184009a --- /dev/null +++ b/test/vision/vgg.jl @@ -0,0 +1,13 @@ +using Boltz, Test + +include("testutils.jl") + +@testset "VGG: $(depth): pretrained: $(pretrained) batchnorm: $(batchnorm)" for depth in [ + 11, 13, 16, 19 + ], + pretrained in [false, true], + batchnorm in [false, true] + + model = Vision.VGG(depth; pretrained, batchnorm) + VisionTestUtils.test_model(model; pretrained) +end diff --git a/test/vision/vit.jl b/test/vision/vit.jl new file mode 100644 index 00000000..6d21c7c3 --- /dev/null +++ b/test/vision/vit.jl @@ -0,0 +1,16 @@ +using Boltz, Test + +include("testutils.jl") + +variants = if parse(Bool, get(ENV, "CI", "false")) + [:tiny, :small, :base] +else + [:tiny, :small, :base, :large, :huge, :giant, :gigantic] +end + +@testset "VisionTransformer: $(variant): pretrained: $(pretrained)" for variant in variants, + pretrained in [false] + + model = Vision.VisionTransformer(variant; pretrained) + VisionTestUtils.test_model(model; pretrained, size=256) +end diff --git a/test/vision/wideresnet.jl b/test/vision/wideresnet.jl new file mode 100644 index 00000000..f5f94752 --- /dev/null +++ b/test/vision/wideresnet.jl @@ -0,0 +1,13 @@ +using Boltz, Test, Metalhead + +include("testutils.jl") + +@testset "WideResNet: $(depth) : pretrained: $(pretrained)" for depth in [50, 101, 152], + pretrained in [false, true] + + depth == 152 && pretrained && continue + pretrained && pkgversion(Metalhead) > v"0.9.4" && continue + + model = Vision.WideResNet(depth; pretrained) + VisionTestUtils.test_model(model; pretrained) +end diff --git a/test/vision_tests.jl b/test/vision_tests.jl deleted file mode 100644 index 27f936a4..00000000 --- a/test/vision_tests.jl +++ /dev/null @@ -1,505 +0,0 @@ -@testsetup module PretrainedWeightsTestSetup - -using Lux, Downloads, JLD2, Pickle, Reactant - -function normalize_imagenet(data) - cmean = reshape(Float32[0.485, 0.456, 0.406], (1, 1, 3, 1)) - cstd = reshape(Float32[0.229, 0.224, 0.225], (1, 1, 3, 1)) - return (data .- cmean) ./ cstd -end - -# The images are normalized and saved -@load joinpath(@__DIR__, "testimages", "monarch_color.jld2") monarch_color_224 monarch_color_256 -const MONARCH_224 = monarch_color_224 -const MONARCH_256 = monarch_color_256 - -const TEST_LBLS = readlines( - Downloads.download( - "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt" - ), -) - -function get_test_image(size, dev) - if size == 224 - return dev(MONARCH_224) - elseif size == 256 - return dev(MONARCH_256) - else - error("size must be 224 or 256") - end -end - -function imagenet_acctest(model, ps, st, dev; size=224) - ps = dev(ps) - st = dev(Lux.testmode(st)) - TEST_X = get_test_image(size, dev) - x = dev(TEST_X) - - if dev isa ReactantDevice - model = Reactant.with_config(; - dot_general_precision=PrecisionConfig.HIGH, - convolution_precision=PrecisionConfig.HIGH, - ) do - @compile model(x, ps, st) - end - end - - ypred = vec(collect(cpu_device()(first(model(x, ps, st))))) - top5 = TEST_LBLS[partialsortperm(ypred, 1:5; rev=true)] - return "monarch" in top5 -end - -export imagenet_acctest, get_test_image - -end - -@testitem "AlexNet" setup = [SharedTestSetup, PretrainedWeightsTestSetup] tags = [:vision] begin - @testset for (mode, aType, dev) in MODES - @testset "pretrained: $(pretrained)" for pretrained in [true, false] - model = Vision.AlexNet(; pretrained) - ps, st = dev(Lux.setup(Random.default_rng(), model)) - st = Lux.testmode(st) - img = get_test_image(224, aType) - - @test size(first(model(img, ps, st))) == (1000, 1) - - if pretrained - @test imagenet_acctest(model, ps, st, dev) - end - - GC.gc(true) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra, img_ra = rdev(cpu_device()((ps, st, img))) - - Reactant.with_config(; - dot_general_precision=PrecisionConfig.HIGH, - convolution_precision=PrecisionConfig.HIGH, - ) do - @test @jit(model(img_ra, ps_ra, st_ra))[1] ≈ model(img, ps, st)[1] atol = - 1e-3 rtol = 1e-3 - end - - if pretrained - @test imagenet_acctest(model, ps, st, rdev) - end - end - end - end -end - -@testitem "ConvMixer" setup = [SharedTestSetup, PretrainedWeightsTestSetup] tags = [ - :vision_metalhead -] begin - using Metalhead: Metalhead - - @testset for (mode, aType, dev) in MODES, name in [:small, :base, :large] - model = Vision.ConvMixer(name; pretrained=false) - ps, st = dev(Lux.setup(Random.default_rng(), model)) - st = Lux.testmode(st) - img = get_test_image(256, aType) - - @test size(first(model(img, ps, st))) == (1000, 1) - - GC.gc(true) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra, img_ra = rdev(cpu_device()((ps, st, img))) - - Reactant.with_config(; - dot_general_precision=PrecisionConfig.HIGH, - convolution_precision=PrecisionConfig.HIGH, - ) do - @test @jit(model(img_ra, ps_ra, st_ra))[1] ≈ model(img, ps, st)[1] atol = - 1e-3 rtol = 1e-3 - end - end - end -end - -@testitem "GoogLeNet" setup = [SharedTestSetup, PretrainedWeightsTestSetup] tags = [ - :vision_metalhead -] begin - using Metalhead: Metalhead - - @testset for (mode, aType, dev) in MODES - model = Vision.GoogLeNet(; pretrained=false) - ps, st = dev(Lux.setup(Random.default_rng(), model)) - st = Lux.testmode(st) - img = get_test_image(224, aType) - - @test size(first(model(img, ps, st))) == (1000, 1) - - GC.gc(true) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra, img_ra = rdev(cpu_device()((ps, st, img))) - - Reactant.with_config(; - dot_general_precision=PrecisionConfig.HIGH, - convolution_precision=PrecisionConfig.HIGH, - ) do - @test @jit(model(img_ra, ps_ra, st_ra))[1] ≈ model(img, ps, st)[1] atol = - 1e-3 rtol = 1e-3 - end - end - end -end - -@testitem "MobileNet" setup = [SharedTestSetup, PretrainedWeightsTestSetup] tags = [ - :vision_metalhead -] begin - using Metalhead: Metalhead - - @testset for (mode, aType, dev) in MODES, name in [:v1, :v2, :v3_small, :v3_large] - model = Vision.MobileNet(name; pretrained=false) - ps, st = dev(Lux.setup(Random.default_rng(), model)) - st = Lux.testmode(st) - img = get_test_image(224, aType) - - @test size(first(model(img, ps, st))) == (1000, 1) - - GC.gc(true) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra, img_ra = rdev(cpu_device()((ps, st, img))) - - Reactant.with_config(; - dot_general_precision=PrecisionConfig.HIGH, - convolution_precision=PrecisionConfig.HIGH, - ) do - @test @jit(model(img_ra, ps_ra, st_ra))[1] ≈ model(img, ps, st)[1] atol = - 1e-3 rtol = 1e-3 - end - end - end -end - -@testitem "ResNet" setup = [SharedTestSetup, PretrainedWeightsTestSetup] tags = [ - :vision_metalhead -] begin - using Metalhead: Metalhead - - @testset for (mode, aType, dev) in MODES - @testset for depth in [18, 34, 50, 101, 152], pretrained in [false, true] - pretrained && pkgversion(Metalhead) > v"0.9.4" && continue - - model = Vision.ResNet(depth; pretrained) - ps, st = dev(Lux.setup(Random.default_rng(), model)) - st = Lux.testmode(st) - img = get_test_image(224, aType) - - @test size(first(model(img, ps, st))) == (1000, 1) - - if pretrained - @test imagenet_acctest(model, ps, st, dev) - end - - GC.gc(true) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra, img_ra = rdev(cpu_device()((ps, st, img))) - - Reactant.with_config(; - dot_general_precision=PrecisionConfig.HIGH, - convolution_precision=PrecisionConfig.HIGH, - ) do - @test @jit(model(img_ra, ps_ra, st_ra))[1] ≈ model(img, ps, st)[1] atol = - 1e-3 rtol = 1e-3 - end - - if pretrained - @test imagenet_acctest(model, ps, st, rdev) - end - end - end - end -end - -@testitem "ResNeXt" setup = [SharedTestSetup, PretrainedWeightsTestSetup] tags = [ - :vision_metalhead -] begin - using Metalhead: Metalhead - - @testset for (mode, aType, dev) in MODES - @testset for (depth, cardinality, base_width) in - [(50, 32, 4), (101, 32, 8), (101, 64, 4), (152, 64, 4)] - @testset for pretrained in [false, true] - depth == 152 && pretrained && continue - pretrained && pkgversion(Metalhead) > v"0.9.4" && continue - - model = Vision.ResNeXt(depth; pretrained, cardinality, base_width) - ps, st = dev(Lux.setup(Random.default_rng(), model)) - st = Lux.testmode(st) - img = get_test_image(224, aType) - - @test size(first(model(img, ps, st))) == (1000, 1) - - if pretrained - @test imagenet_acctest(model, ps, st, dev) - end - - GC.gc(true) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra, img_ra = rdev(cpu_device()((ps, st, img))) - - Reactant.with_config(; - dot_general_precision=PrecisionConfig.HIGH, - convolution_precision=PrecisionConfig.HIGH, - ) do - @test @jit(model(img_ra, ps_ra, st_ra))[1] ≈ model(img, ps, st)[1] atol = - 1e-3 rtol = 1e-3 - end - - if pretrained - @test imagenet_acctest(model, ps, st, rdev) - end - end - end - end - end -end - -@testitem "WideResNet" setup = [SharedTestSetup, PretrainedWeightsTestSetup] tags = [ - :vision_metalhead -] begin - using Metalhead: Metalhead - - @testset for (mode, aType, dev) in MODES - @testset for depth in [50, 101, 152], pretrained in [false, true] - depth == 152 && pretrained && continue - pretrained && pkgversion(Metalhead) > v"0.9.4" && continue - - model = Vision.WideResNet(depth; pretrained) - ps, st = dev(Lux.setup(Random.default_rng(), model)) - st = Lux.testmode(st) - img = get_test_image(224, aType) - - @test size(first(model(img, ps, st))) == (1000, 1) - - if pretrained - @test imagenet_acctest(model, ps, st, dev) - end - - GC.gc(true) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra, img_ra = rdev(cpu_device()((ps, st, img))) - - Reactant.with_config(; - dot_general_precision=PrecisionConfig.HIGH, - convolution_precision=PrecisionConfig.HIGH, - ) do - @test @jit(model(img_ra, ps_ra, st_ra))[1] ≈ model(img, ps, st)[1] atol = - 1e-3 rtol = 1e-3 - end - end - end - end -end - -@testitem "SqueezeNet" setup = [SharedTestSetup, PretrainedWeightsTestSetup] tags = [ - :vision_metalhead -] begin - using Metalhead: Metalhead - - @testset for (mode, aType, dev) in MODES - @testset for pretrained in [false, true] - model = Vision.SqueezeNet(; pretrained) - ps, st = dev(Lux.setup(Random.default_rng(), model)) - st = Lux.testmode(st) - img = get_test_image(224, aType) - - @test size(first(model(img, ps, st))) == (1000, 1) - - if pretrained - @test imagenet_acctest(model, ps, st, dev) - end - - GC.gc(true) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra, img_ra = rdev(cpu_device()((ps, st, img))) - - Reactant.with_config(; - dot_general_precision=PrecisionConfig.HIGH, - convolution_precision=PrecisionConfig.HIGH, - ) do - @test @jit(model(img_ra, ps_ra, st_ra))[1] ≈ model(img, ps, st)[1] atol = - 1e-3 rtol = 1e-3 - end - - if pretrained - @test imagenet_acctest(model, ps, st, rdev) - end - end - end - end -end - -@testitem "VGG" setup = [SharedTestSetup, PretrainedWeightsTestSetup] tags = [:vision] begin - @testset for (mode, aType, dev) in MODES - @testset for depth in [11, 13, 16, 19], - pretrained in [false, true], - batchnorm in [false, true] - - depth ≥ 16 && get(ENV, "CI", "false") == "true" && continue - - model = Vision.VGG(depth; batchnorm, pretrained) - ps, st = dev(Lux.setup(Random.default_rng(), model)) - st = Lux.testmode(st) - img = get_test_image(224, aType) - - @test size(first(model(img, ps, st))) == (1000, 1) - - if pretrained - @test imagenet_acctest(model, ps, st, dev) - end - - GC.gc(true) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra, img_ra = rdev(cpu_device()((ps, st, img))) - - Reactant.with_config(; - dot_general_precision=PrecisionConfig.HIGH, - convolution_precision=PrecisionConfig.HIGH, - ) do - @test @jit(model(img_ra, ps_ra, st_ra))[1] ≈ model(img, ps, st)[1] atol = - 1e-3 rtol = 1e-3 - end - - if pretrained - @test imagenet_acctest(model, ps, st, rdev) - end - end - end - end -end - -@testitem "EfficientNet" setup = [SharedTestSetup, PretrainedWeightsTestSetup] tags = [ - :vision -] begin - all_names = if parse(Bool, get(ENV, "CI", "false")) - [:b0, :b1, :b2] - else - [:b0, :b1, :b2, :b3, :b4, :b5, :b6, :b7] - end - @testset for (mode, aType, dev) in MODES - @testset for name in all_names, pretrained in [false, true] - model = Boltz.Vision.EfficientNet(name; pretrained) - ps, st = Lux.setup(Random.default_rng(), model) |> dev - st = Lux.testmode(st) - img = get_test_image(224, aType) - - @test size(first(model(img, ps, st))) == (1000, 1) - - if pretrained - @test imagenet_acctest(model, ps, st, dev) - end - - GC.gc(true) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra, img_ra = rdev(cpu_device()((ps, st, img))) - - Reactant.with_config(; - dot_general_precision=PrecisionConfig.HIGH, - convolution_precision=PrecisionConfig.HIGH, - ) do - @test @jit(model(img_ra, ps_ra, st_ra))[1] ≈ model(img, ps, st)[1] atol = - 1e-3 rtol = 1e-3 - end - - if pretrained - @test imagenet_acctest(model, ps, st, rdev) - end - end - end - end -end - -@testitem "VisionTransformer" setup = [SharedTestSetup, PretrainedWeightsTestSetup] tags = [ - :vision -] begin - all_names = if parse(Bool, get(ENV, "CI", "false")) - [:tiny, :small, :base] - else - [:tiny, :small, :base, :large, :huge, :giant, :gigantic] - end - - @testset for (mode, aType, dev) in MODES, name in all_names - model = Vision.VisionTransformer(name; pretrained=false) - ps, st = dev(Lux.setup(Random.default_rng(), model)) - st = Lux.testmode(st) - img = get_test_image(256, aType) - - @test size(first(model(img, ps, st))) == (1000, 1) - - model = Vision.VisionTransformer(name; pretrained=false) - ps, st = dev(Lux.setup(Random.default_rng(), model)) - st = Lux.testmode(st) - img = aType(randn(Float32, 256, 256, 3, 2)) - - @test size(first(model(img, ps, st))) == (1000, 2) - - GC.gc(true) - - if test_reactant(mode) - set_reactant_backend!(mode) - - rdev = reactant_device(; force=true) - - ps_ra, st_ra, img_ra = rdev(cpu_device()((ps, st, img))) - - Reactant.with_config(; - dot_general_precision=PrecisionConfig.HIGH, - convolution_precision=PrecisionConfig.HIGH, - ) do - @test @jit(model(img_ra, ps_ra, st_ra))[1] ≈ model(img, ps, st)[1] atol = - 1e-3 rtol = 1e-3 - end - end - end -end