diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml new file mode 100644 index 0000000..5c56b17 --- /dev/null +++ b/.github/workflows/benchmark.yml @@ -0,0 +1,29 @@ +name: Benchmark this PR + +on: + pull_request: + branches: [main] + paths-ignore: ['docs/**'] + +permissions: + pull-requests: write + contents: read + +jobs: + benchmark: + runs-on: ubuntu-latest + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + version: ["1", "lts"] + env: + # Force consistent Julia depot path for self-hosted runners + JULIA_DEPOT_PATH: ~/.julia + steps: + - uses: MilesCranmer/AirspeedVelocity.jl@action-v1 + with: + julia-version: ${{ matrix.version }} + script: "benchmark/benchmarks.jl" + extra-pkgs: "Lux,StableRNGs" + job-summary: true diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl new file mode 100644 index 0000000..5a52f79 --- /dev/null +++ b/benchmark/benchmarks.jl @@ -0,0 +1,43 @@ +using NeuralOperators, BenchmarkTools +using Lux, Random, StableRNGs + +const SUITE = BenchmarkGroup() +const rng = StableRNG(123) + +# ============================================================================= +# FourierNeuralOperator +# ============================================================================= + +fno1 = FourierNeuralOperator(; chs = (2, 8, 8, 8, 4), modes = (8,), shift = false) +fno2 = FourierNeuralOperator(; chs = (2, 8, 8, 8, 4), modes = (4, 4), shift = true) + +ps1, st1 = Lux.setup(rng, fno1) +ps2, st2 = Lux.setup(rng, fno2) + +x1 = rand(rng, Float32, 32, 2, 4) +x2 = rand(rng, Float32, 16, 16, 2, 4) + +SUITE["fno"] = BenchmarkGroup() + +SUITE["fno"]["construct_1d"] = @benchmarkable FourierNeuralOperator( + ; chs = (2, 8, 8, 8, 4), modes = (8,) +) +SUITE["fno"]["setup_1d"] = @benchmarkable Lux.setup($rng, $fno1) +SUITE["fno"]["forward_1d"] = @benchmarkable $fno1($x1, $ps1, $st1) +SUITE["fno"]["forward_2d"] = @benchmarkable $fno2($x2, $ps2, $st2) + +# ============================================================================= +# DeepONet +# ============================================================================= + +don = DeepONet(; branch = (64, 32, 32, 16), trunk = (1, 8, 8, 16)) +psd, std = Lux.setup(rng, don) + +xin = (rand(rng, Float32, 64, 5), rand(rng, Float32, 1, 10)) + +SUITE["deeponet"] = BenchmarkGroup() + +SUITE["deeponet"]["construct"] = @benchmarkable DeepONet( + ; branch = (64, 32, 32, 16), trunk = (1, 8, 8, 16) +) +SUITE["deeponet"]["forward"] = @benchmarkable $don($xin, $psd, $std)