Skip to content

Neuroevolution on the GPU

The problem

This example is about the engine: it evaluates a whole generation on the GPU. The task is to fit a small neural network to a function of two variables, sin(2x) cos(y), sampled at 4,096 points of a 64 × 64 grid over [−2, 2]². The network has 2 inputs, 16 hidden tanh units and a linear output:

output = c + Σ vⱼ tanh(aⱼ x + bⱼ y + dⱼ)    over the 16 hidden units j

The fitness is the mean squared error between the network's output and the function, over the 4,096 samples.

There's no Python version: the example is a Rust crate of its own, with wgpu. In Python, batch=True hands a whole generation to the fitness function, which can pass it to any GPU library.

What makes it hard

The cost of the fitness. One evaluation runs the network on 4,096 samples. A generation of 512 networks takes about 2 million network evaluations. The work is the same for every network and every sample, and independent: that suits a GPU, which runs thousands of such computations at once.

The search. The 65 weights interact: an output weight only matters through its hidden unit's input weights and bias, and swapping two hidden units gives the same network. A good fit takes hundreds of thousands of evaluations.

Representation

A Real genome of 65 genes in [−3, 3]: per hidden unit, its two input weights and its bias (48 genes), then the 16 output weights, then the output bias. The network is decoded from the genome for every evaluation, on the CPU or on the GPU.

Algorithm

The main run is CMA-ES (Hansen and Ostermeier, 2001, Evolutionary Computation 9(2): 159-195), which genoxide's guide recommends for continuous problems of up to a few hundred genes whose genes interact. It samples a population from a normal distribution, and adapts its mean, step size and covariance matrix to the steps that worked, which suits the correlated weights of a network. It has a population of 128 instead of the default 4 + ⌊3 ln 65⌋ = 16, since a batch of 16 networks leaves most of the GPU idle: with 16, seed 1 took about as many evaluations, 249,000, but 9.5 s, most of them spent on the round trips of small batches. The initial step size is genoxide's default, 0.3 of each gene's range. IPOP restarts (Auger and Hansen, 2005) start again from a random point with twice the population if a run converges before the target. The run stops at an error of 10⁻⁵, or after 2,000,000 evaluations.

Every generation goes to the GPU through Batch, which scores a whole generation in one call. The call uploads the weights, dispatches a WGSL compute shader with wgpu, and downloads the errors. Each genome gets a workgroup of 256 threads; each thread takes a share of the samples, and the workgroup adds their errors up. The GPU computes in single precision.

Then, as a contrast, a genetic algorithm runs 300 generations of 512 networks, with tournaments of size 3, simulated binary crossover with η = 15, and polynomial mutation with η = 20 at a rate of 2/65 per gene, two genes per child on average. It runs twice, with the same seed:

  • on the GPU, through the same Batch;
  • on the CPU, with Engine::parallel(true): rayon spreads the generation's evaluations over every core, in double precision.

The two runs compare the speed of the evaluations; the genetic algorithm isn't meant to fit the network well.

The example is a crate of its own, since it depends on wgpu, and is Rust only. Run it with:

cargo run --release --manifest-path examples/gpu/Cargo.toml

Output

The first line names the GPU; without one, the example says so and stops. The second gives the size of the task. Then there is a line per run: its time, its evaluations, and the best network's error. Under the CMA-ES line, that network's error is recomputed on the CPU in double precision, which shows how much single precision changed it.

CI only compiles this example, since it has no GPU: output.txt comes from a run on an NVIDIA GeForce RTX 4060, with 20 CPU threads. The times depend on the machine, and change from run to run.

The project page plays back the CMA-ES run of output.txt.

Good results

The function isn't known to be exactly representable by this network, so its best fit isn't known. For scale, the samples have a variance of 0.177: a network that outputs 0 everywhere has that error. Gradient descent, which genoxide doesn't do, gives a reference: L-BFGS-B with the weights in the same [−3, 3], from 6 random starts and up to 100,000 iterations each, reached errors of 4.4 × 10⁻⁹ to 1.2 × 10⁻⁶, the best 4.4 × 10⁻⁹, in several minutes each.

CMA-ES gets to 10⁻⁵, about 0.006% of the variance: in the run of output.txt, after 296,704 evaluations and 2.90 s, with the same error in double precision. Over seeds 1 to 20 on the RTX 4060, every run got there, after 243,000 to 640,000 evaluations (289,000 at the median), in 2.3 to 6 seconds. Getting further is slow: two runs aimed at 10⁻⁶ got there after 575,000 and 1,600,000 evaluations, 7 and 18 seconds. The best fits of gradient descent, a thousand times smaller than 10⁻⁶, are out of reach of CMA-ES in seconds: it doesn't reach the best fit known.

The genetic algorithm stops at 0.0115 after its 300 generations, over a thousand times the error of CMA-ES with half the evaluations. On the GPU it took 0.55 s, against 4.35 s on every CPU core, eight times as long.

Known optimum: none known; the best fit found by gradient descent: 4.4e-9 (mean squared error)

Source: examples/gpu

Interactive run: tachsin.gr/projects/genoxide/examples/gpu

cargo run --release --manifest-path examples/gpu/Cargo.toml

//! Batch fitness evaluation on the GPU: neuroevolution of a small neural network. Each genome is
//! the 65 weights of a network with 2 inputs, 16 hidden tanh units and 1 output. Its fitness is
//! the mean squared error of the network over 4,096 samples of a function: about 2 million
//! network evaluations per generation. A generation goes to the GPU at once, through `Batch`: one
//! upload of the weights, one dispatch with a workgroup per genome, one download of the errors.
//!
//! CMA-ES with IPOP restarts fits the network to an error of 1e-5 on the GPU. A genetic algorithm
//! then runs 300 generations of 512 networks twice, on the GPU and on every CPU core, to compare
//! their speed; it stops far from a good fit.
//!
//! With `GENOXIDE_TRACE=<file>`, it also writes a trace of its run for the plot on the example's
//! page, with `trace.rs`.
//!
//! cargo run --release --manifest-path examples/gpu/Cargo.toml

mod trace;

use genoxide::prelude::*;
use std::sync::Mutex;
use std::time::Instant;
use wgpu::util::DeviceExt;

const HIDDEN: usize = 16;
// per hidden unit: a weight for each input and a bias; then the output weights and bias
const WEIGHTS: usize = HIDDEN * 3 + HIDDEN + 1;
const SAMPLES: usize = 4_096;
// the error CMA-ES stops at
const TARGET: f64 = 1e-5;
// the genetic algorithm's population and generations
const POPULATION: usize = 512;
const GENERATIONS: u64 = 300;

// a workgroup per genome: its 256 invocations share the samples, then add up their errors
const SHADER: &str = "
const HIDDEN: u32 = 16u;
const WEIGHTS: u32 = 65u;

@group(0) @binding(0) var<storage, read> weights: array<f32>;
@group(0) @binding(1) var<storage, read> samples: array<vec4<f32>>; // x, y, target, unused
@group(0) @binding(2) var<storage, read_write> errors: array<f32>;

var<workgroup> partial: array<f32, 256>;

@compute @workgroup_size(256)
fn main(@builtin(workgroup_id) group: vec3<u32>, @builtin(local_invocation_index) local: u32) {
    let w = group.x * WEIGHTS;
    let count = arrayLength(&samples);
    var error = 0.0;
    for (var s = local; s < count; s = s + 256u) {
        let sample = samples[s];
        var output = weights[w + HIDDEN * 4u];
        for (var h = 0u; h < HIDDEN; h = h + 1u) {
            let unit = w + h * 3u;
            let activation = weights[unit] * sample.x + weights[unit + 1u] * sample.y + weights[unit + 2u];
            output = output + weights[w + HIDDEN * 3u + h] * tanh(activation);
        }
        let difference = output - sample.z;
        error = error + difference * difference;
    }
    partial[local] = error;
    workgroupBarrier();
    for (var stride = 128u; stride > 0u; stride = stride / 2u) {
        if (local < stride) {
            partial[local] = partial[local] + partial[local + stride];
        }
        workgroupBarrier();
    }
    if (local == 0u) {
        errors[group.x] = partial[0] / f32(count);
    }
}
";

/// The samples: a grid over [-2, 2]², with the target sin(2x)·cos(y).
fn samples() -> Vec<[f64; 3]> {
    let side = (SAMPLES as f64).sqrt() as usize;
    let coordinate = |i: usize| -2.0 + 4.0 * i as f64 / (side - 1) as f64;
    (0..side * side)
        .map(|i| {
            let (x, y) = (coordinate(i % side), coordinate(i / side));
            [x, y, (2.0 * x).sin() * y.cos()]
        })
        .collect()
}

/// The network's mean squared error over the samples, on the CPU.
fn error(weights: &Reals, samples: &[[f64; 3]]) -> f64 {
    let total: f64 = samples
        .iter()
        .map(|&[x, y, target]| {
            let mut output = weights[HIDDEN * 4];
            for h in 0..HIDDEN {
                let unit = h * 3;
                let activation = weights[unit] * x + weights[unit + 1] * y + weights[unit + 2];
                output += weights[HIDDEN * 3 + h] * activation.tanh();
            }
            (output - target).powi(2)
        })
        .sum();
    total / samples.len() as f64
}

/// The network's error on the GPU, in single precision.
struct GpuError {
    device: wgpu::Device,
    queue: wgpu::Queue,
    pipeline: wgpu::ComputePipeline,
    samples: wgpu::Buffer,
    // the weights of a generation, to upload; reused
    staging: Mutex<Vec<f32>>,
}

impl GpuError {
    /// The GPU's device, with the samples uploaded, or `None` without a GPU.
    fn new(samples: &[[f64; 3]]) -> Option<Self> {
        let instance = wgpu::Instance::default();
        let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
            power_preference: wgpu::PowerPreference::HighPerformance,
            ..Default::default()
        }))
        .ok()?;
        println!("GPU: {}", adapter.get_info().name);
        let (device, queue) =
            pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor::default())).ok()?;
        let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("error"),
            source: wgpu::ShaderSource::Wgsl(SHADER.into()),
        });
        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("error"),
            layout: None,
            module: &module,
            entry_point: Some("main"),
            compilation_options: Default::default(),
            cache: None,
        });
        let packed: Vec<[f32; 4]> = samples
            .iter()
            .map(|&[x, y, target]| [x as f32, y as f32, target as f32, 0.0])
            .collect();
        let samples = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("samples"),
            contents: bytemuck::cast_slice(&packed),
            usage: wgpu::BufferUsages::STORAGE,
        });
        Some(Self {
            device,
            queue,
            pipeline,
            samples,
            staging: Mutex::new(Vec::new()),
        })
    }

    /// The errors of the networks `genomes`, in their order.
    fn evaluate(&self, genomes: &[&Reals]) -> Vec<f64> {
        if genomes.is_empty() {
            return Vec::new();
        }
        let mut staging = self.staging.lock().unwrap();
        staging.clear();
        for genome in genomes {
            staging.extend(genome.iter().map(|&weight| weight as f32));
        }
        let weights = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("weights"),
                contents: bytemuck::cast_slice(&staging),
                usage: wgpu::BufferUsages::STORAGE,
            });
        drop(staging);
        let size = (genomes.len() * size_of::<f32>()) as u64;
        let errors = self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("errors"),
            size,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        let download = self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("download"),
            size,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let bindings = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &self.pipeline.get_bind_group_layout(0),
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: weights.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: self.samples.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: errors.as_entire_binding(),
                },
            ],
        });
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
            pass.set_pipeline(&self.pipeline);
            pass.set_bind_group(0, &bindings, &[]);
            pass.dispatch_workgroups(genomes.len() as u32, 1, 1);
        }
        encoder.copy_buffer_to_buffer(&errors, 0, &download, 0, size);
        self.queue.submit([encoder.finish()]);

        let slice = download.slice(..);
        slice.map_async(wgpu::MapMode::Read, |result| {
            result.expect("the GPU reads back the errors");
        });
        self.device
            .poll(wgpu::PollType::wait_indefinitely())
            .expect("the GPU finishes");
        let values: Vec<f64> = bytemuck::cast_slice::<u8, f32>(
            &slice.get_mapped_range().expect("the errors are mapped"),
        )
        .iter()
        .map(|&error| f64::from(error))
        .collect();
        download.unmap();
        values
    }
}

fn cmaes() -> genoxide::Result<Cmaes> {
    Cmaes::builder(Real::uniform(WEIGHTS, -3.0..=3.0)?)
        .population_size(128)
        .restarts(cmaes::Restarts::Ipop)
        .minimize()
        .seed(1)
        .build()
}

fn ga() -> genoxide::Result<Ga<Real, Tournament, SimulatedBinaryCrossover, PolynomialMutation>> {
    Ga::builder(Real::uniform(WEIGHTS, -3.0..=3.0)?)
        .population_size(POPULATION)
        .select(Tournament::new(3)?)
        .crossover(SimulatedBinaryCrossover::new(15.0)?)
        .mutate(PolynomialMutation::per_gene(2.0 / WEIGHTS as f64, 20.0)?)
        .minimize()
        .seed(1)
        .build()
}

fn main() -> genoxide::Result<()> {
    let samples = samples();
    let Some(gpu) = GpuError::new(&samples) else {
        println!("no GPU found: nothing to compare");
        return Ok(());
    };
    println!(
        "a {WEIGHTS}-weight network fitted to {SAMPLES} samples
"
    );

    // with GENOXIDE_TRACE=<file>, a trace of the run for the plot on the example's page
    let mut trace = trace::Trace::from_env();
    let start = Instant::now();
    let evaluate = trace.timed(|genomes: &[&Reals]| gpu.evaluate(genomes));
    let outcome = Engine::new(cmaes()?, Batch(evaluate))
        .stop_when(Stop::target(TARGET).or(Stop::evaluations(2_000_000)))
        .on_generation(|snapshot| trace.record(snapshot))
        .run()?;
    report("CMA-ES on the GPU", &outcome, start.elapsed());
    let double = error(outcome.best_genome(), &samples);
    println!(
        "{:19}in double precision on the CPU: error {double:.2e}",
        ""
    );
    trace.write();

    println!(
        "
a genetic algorithm, {GENERATIONS} generations of {POPULATION} networks:"
    );
    let start = Instant::now();
    let outcome = Engine::new(ga()?, Batch(|genomes: &[&Reals]| gpu.evaluate(genomes)))
        .stop_when(Stop::generations(GENERATIONS))
        .run()?;
    report("on the GPU", &outcome, start.elapsed());
    let start = Instant::now();
    let outcome = Engine::new(ga()?, |weights: &Reals| error(weights, &samples))
        .parallel(true)
        .stop_when(Stop::generations(GENERATIONS))
        .run()?;
    report("on every CPU core", &outcome, start.elapsed());
    Ok(())
}

fn report(name: &str, outcome: &Outcome<Reals>, elapsed: std::time::Duration) {
    println!(
        "{name:<19}{:>5.2} s, {:>7} evaluations, error {:.2e}",
        elapsed.as_secs_f64(),
        outcome.evaluations(),
        outcome.best_fitness().score().unwrap_or(f64::NAN)
    );
}

What it prints, from a seeded run:

GPU: NVIDIA GeForce RTX 4060
a 65-weight network fitted to 4096 samples

CMA-ES on the GPU   2.90 s,  296704 evaluations, error 9.97e-6
                   in double precision on the CPU: error 9.97e-6

a genetic algorithm, 300 generations of 512 networks:
on the GPU          0.55 s,  151670 evaluations, error 1.15e-2
on every CPU core   4.35 s,  151670 evaluations, error 1.15e-2