Hi, I am training a Neural Operator to learn a map between a two dimensional domain. For Julia, this is a MWE of the problem I am considering:
# Imports. using NeuralOperators using Flux using FluxTraining using MLUtils # Generate data. N_samples = 5_000 xdata = Float32.(rand(1, 16, 16, N_samples)) ydata = Float32.(rand(1, 16, 16, N_samples)) # Model parameters. optimiser = Flux.Optimiser(WeightDecay(1.0f-4), Flux.Adam(5.0f-4)) # DataLoaders. data_train, data_test = splitobs((xdata, ydata), at = 0.9) loader_train, loader_test = DataLoader(data_train, batchsize = 100), DataLoader(data_test, batchsize = 100) data = collect.((loader_train, loader_test)) chs = (1, 32, 32, 32, 32, 32, 64, 1) model = FourierNeuralOperator(ch = chs, modes = (8, 8), σ = gelu) learner = Learner(model, data, optimiser, l₂loss) for k in range(1, 100) epoch!(learner, TrainingPhase(), learner.data.training) end
It takes ~30-40s to train an epoch (even when using
julia --threads=auto). In Python, the equivalent problem takes almost an order of magnitude less time while maintaining the same error by epoch.
How can I get closer to Python’s performance in Julia?
Any help will be appreciated. Thanks!