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!