From eba6f2643787178253bbbb36da64c8d54d20d365 Mon Sep 17 00:00:00 2001 From: Jonathan Merritt Date: Tue, 24 Sep 2024 17:21:14 +1000 Subject: [PATCH] GitHub CI workflows --- .github/workflows/dry-run.yml | 45 ++++++++++++++++++++++++++++++++ README.md | 2 ++ generate_data/flake.nix | 17 +++++++----- generate_data/quadcopter.py | 4 +-- nn/.github/workflows/dry-run.yml | 18 ------------- nn/Cargo.toml | 2 +- nn/flake.nix | 39 ++++++++++++++++++--------- nn/src/bin/main_train.rs | 12 +++++++-- nn/src/bin/main_vis_results.rs | 19 +++++++++++--- nn/src/training.rs | 31 +++++++++++++++++++--- walkthrough.md | 12 +++++++++ 11 files changed, 151 insertions(+), 50 deletions(-) create mode 100644 .github/workflows/dry-run.yml delete mode 100644 nn/.github/workflows/dry-run.yml diff --git a/.github/workflows/dry-run.yml b/.github/workflows/dry-run.yml new file mode 100644 index 0000000..3255c3d --- /dev/null +++ b/.github/workflows/dry-run.yml @@ -0,0 +1,45 @@ +name: "Dry Run" +on: + pull_request: + push: +jobs: + run_steps_from_readme: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/cache@v3 + with: + path: | + nn/target + key: ${{ runner.os }}-cache + - uses: cachix/install-nix-action@v27 + with: + nix_path: nixpkgs=channel:nixos-unstable + + - name: Compile Rust Code + run: | + echo "Compiling Rust Code" + pushd nn + nix develop . --command cargo build --release + popd + + - name: Generate Optimized Trajectories + run: | + echo "Generating Optimised Trajectories" + pushd generate_data + nix develop . --command env GENERATE_DATA_N_TRAJECTORIES=100 uv run quadcopter.py + popd + + - name: Run Training + run: | + echo "Running Training" + pushd nn + nix develop . --command env TRAIN_USE_NDARRAY=1 cargo run --release --bin train + popd + + - name: Run Inference + run: | + echo "Running Inference" + pushd nn + nix develop . --command env INFERENCE_USE_NDARRAY=1 cargo run --release --bin vis_results + popd \ No newline at end of file diff --git a/README.md b/README.md index 6b93e70..9ec69ad 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,7 @@ # Quadcopter Optimal Landing: Neural Network Controller +![GitHub CI](https://github.com/lancelet/nn-landing-poc/actions/workflows/dry-run.yml/badge.svg) + WARNING: This is an initial proof-of-concept, not a carefully polished project! diff --git a/generate_data/flake.nix b/generate_data/flake.nix index b035db0..27d5251 100644 --- a/generate_data/flake.nix +++ b/generate_data/flake.nix @@ -6,17 +6,22 @@ flake-utils.url = "github:numtide/flake-utils"; }; - outputs = { self, nixpkgs, flake-utils }: - flake-utils.lib.eachDefaultSystem (system: - let - pkgs = import nixpkgs { inherit system; }; - in - { + outputs = { + self, + nixpkgs, + flake-utils, + }: + flake-utils.lib.eachDefaultSystem ( + system: let + pkgs = import nixpkgs {inherit system;}; + in { devShells.default = pkgs.mkShell { buildInputs = with pkgs; [ + stdenv python312 uv ]; + LD_LIBRARY_PATH = "${pkgs.stdenv.cc.cc.lib}/lib"; }; } ); diff --git a/generate_data/quadcopter.py b/generate_data/quadcopter.py index 0cd403d..65fabc0 100644 --- a/generate_data/quadcopter.py +++ b/generate_data/quadcopter.py @@ -21,10 +21,10 @@ #---- Constants --------------------------------------------------------------- # Number of processes to use for data generation. -N_PROCESSES = 8 +N_PROCESSES = int(os.getenv("GENERATE_DATA_N_PROCESSES", 8)) # Number of trajectories to generate. -N_TRAJECTORIES = 15000 +N_TRAJECTORIES = int(os.getenv("GENERATE_DATA_N_TRAJECTORIES", 15000)) # Current project directory (based on the location of this file). PROJECT_DIR = Path(os.path.dirname(os.path.abspath(__file__))).resolve() diff --git a/nn/.github/workflows/dry-run.yml b/nn/.github/workflows/dry-run.yml deleted file mode 100644 index 992a2d6..0000000 --- a/nn/.github/workflows/dry-run.yml +++ /dev/null @@ -1,18 +0,0 @@ -name: "Dry Run" -on: - pull_request: - push: -jobs: - run_steps_from_readme: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - uses: cachix/install-nix-actions@v27 - with: - nix_path: nixpkgs=channel:nixos-unstable - - - name: Compile Rust code - run: | - echo "Compiling Rust code" - pushd nn - nix develop . --command cargo build --release \ No newline at end of file diff --git a/nn/Cargo.toml b/nn/Cargo.toml index ca61c00..dd1e611 100644 --- a/nn/Cargo.toml +++ b/nn/Cargo.toml @@ -5,7 +5,7 @@ edition = "2021" [dependencies] anyhow = "1.0.89" -burn = { version = "0.14.0", features = ["dataset", "train", "wgpu"] } +burn = { version = "0.14.0", features = ["dataset", "train", "wgpu", "ndarray"] } ndarray = "0.16.1" ndarray-npy = "0.9.1" ode_solvers = "0.4.0" diff --git a/nn/flake.nix b/nn/flake.nix index e78f621..454570e 100644 --- a/nn/flake.nix +++ b/nn/flake.nix @@ -15,23 +15,36 @@ system: let pkgs = import nixpkgs {inherit system;}; naersk-lib = pkgs.callPackage naersk {}; + isMacOS = system == "x86_64-darwin" || system == "aarch64-darwin"; + isLinux = system == "x86_64-linux" || system == "aarch64-linux" || system == "i686-linux"; in { defaultPackage = naersk-lib.buildPackage ./.; devShell = with pkgs; mkShell { - buildInputs = [ - cargo - darwin.apple_sdk.frameworks.IOKit - darwin.apple_sdk.frameworks.QuartzCore - ffmpeg - iconv - pre-commit - python3 - python3Packages.matplotlib - rustc - rustfmt - rustPackages.clippy - ]; + buildInputs = + [ + cargo + ffmpeg + iconv + pre-commit + python3 + python3Packages.matplotlib + rustc + rustfmt + rustPackages.clippy + ] + ++ lib.optionals isMacOS [ + darwin.apple_sdk.frameworks.IOKit + darwin.apple_sdk.frameworks.QuartzCore + ] + ++ lib.optionals isLinux [ + vulkan-loader + vulkan-headers + vulkan-validation-layers + vulkan-tools + glslang + shaderc + ]; RUST_SRC_PATH = rustPlatform.rustLibSrc; }; } diff --git a/nn/src/bin/main_train.rs b/nn/src/bin/main_train.rs index d80825b..a3b2df6 100644 --- a/nn/src/bin/main_train.rs +++ b/nn/src/bin/main_train.rs @@ -1,12 +1,20 @@ +use std::env; + use train_quadcopter::training; use anyhow::Result; +use burn::backend::ndarray::{NdArray, NdArrayDevice}; use burn::backend::wgpu::{Wgpu, WgpuDevice}; use burn::backend::Autodiff; fn main() -> Result<()> { println!("Quadcopter Training"); - let device = WgpuDevice::default(); - training::run::>(device) + if env::var("TRAIN_USE_NDARRAY").is_ok() { + let device = NdArrayDevice::default(); + training::run::>(device) + } else { + let device = WgpuDevice::default(); + training::run::>(device) + } } diff --git a/nn/src/bin/main_vis_results.rs b/nn/src/bin/main_vis_results.rs index 14c08e7..f0bc278 100644 --- a/nn/src/bin/main_vis_results.rs +++ b/nn/src/bin/main_vis_results.rs @@ -1,10 +1,13 @@ +use std::env; use std::f32::consts::PI; use std::fs::create_dir_all; use anyhow::{bail, Result}; +use burn::backend::ndarray::NdArrayDevice; use burn::backend::wgpu::WgpuDevice; -use burn::backend::Wgpu; +use burn::backend::{NdArray, Wgpu}; use burn::module::Module; +use burn::prelude::Backend; use burn::record::{DefaultFileRecorder, FullPrecisionSettings}; use plotpy::{Curve, Plot}; @@ -21,10 +24,18 @@ fn main() -> Result<()> { println!("Trained Model Visualization"); // Load the model weights from the trained file. - let device = WgpuDevice::default(); + if env::var("INFERENCE_USE_NDARRAY").is_ok() { + let device = NdArrayDevice::default(); + run_inference::(device) + } else { + let device = WgpuDevice::default(); + run_inference::(device) + } +} + +fn run_inference(device: B::Device) -> Result<()> { let file_path = "./burn-training-artifacts/model.mpk"; - let config = ModelConfig::new(); - let model = config.init::(&device); + let model = ModelConfig::new().init::(&device); let model = model.load_file( file_path, &DefaultFileRecorder::::new(), diff --git a/nn/src/training.rs b/nn/src/training.rs index 8ba243b..2ea044e 100644 --- a/nn/src/training.rs +++ b/nn/src/training.rs @@ -1,3 +1,5 @@ +use std::io::IsTerminal; + use crate::{ data::{load_datasets, ExampleBatcher, MeanStdev, StateMeanStdev}, model::ModelConfig, @@ -14,6 +16,7 @@ use burn::{ store::{Aggregate, Direction, Split}, LossMetric, }, + renderer::MetricsRenderer, LearnerBuilder, MetricEarlyStoppingStrategy, StoppingCondition, }, }; @@ -65,7 +68,7 @@ pub fn run(device: B::Device) -> Result<()> { .build(test_dataset); // Model - let learner = LearnerBuilder::new(ARTIFACT_DIR) + let learner_builder = LearnerBuilder::new(ARTIFACT_DIR) .metric_train_numeric(LossMetric::new()) .metric_valid_numeric(LossMetric::new()) .with_file_checkpointer(DefaultFileRecorder::::new()) @@ -76,9 +79,13 @@ pub fn run(device: B::Device) -> Result<()> { StoppingCondition::NoImprovementSince { n_epochs: 1 }, )) .devices(vec![device.clone()]) - .num_epochs(num_epochs) - .summary() - .build(model, optimizer.init(), lr); + .num_epochs(num_epochs); + let learner_builder = if !std::io::stdout().is_terminal() { + learner_builder.renderer(StdoutMetricsRenderer) + } else { + learner_builder + }; + let learner = learner_builder.summary().build(model, optimizer.init(), lr); let model_trained = learner.fit(dataloader_train, dataloader_test); @@ -91,3 +98,19 @@ pub fn run(device: B::Device) -> Result<()> { Ok(()) } + +struct StdoutMetricsRenderer; + +impl MetricsRenderer for StdoutMetricsRenderer { + fn update_train(&mut self, _state: burn::train::renderer::MetricState) {} + + fn update_valid(&mut self, _state: burn::train::renderer::MetricState) {} + + fn render_train(&mut self, item: burn::train::renderer::TrainingProgress) { + println!("Training iteration: {}", item.iteration); + } + + fn render_valid(&mut self, item: burn::train::renderer::TrainingProgress) { + println!("Validation iteration: {}", item.iteration); + } +} diff --git a/walkthrough.md b/walkthrough.md index 9f40ef9..df2520a 100644 --- a/walkthrough.md +++ b/walkthrough.md @@ -117,3 +117,15 @@ to a movie with `ffmpeg` like this: ```bash ffmpeg -framerate 60 -i ./animation/%05d.png -vcodec libx264 -s 1080x1080 -pix_fmt yuv420p animation.mov ``` + +## CI Build + +There is a CI build which runs the project through its basic steps. See +`.github/workflows/dry-run.yml`. The main differences between the CI build +and running the project locally are: + - The CI build fetches some dependencies using Nix. + - Only 100 training trajectories are generated. + - The renderer for the fancy training GUI is replaced by one that dumps + progress to `stdout`. + - The `NDArray` backend is used instead of `wgpu`. I couldn't figure out + the trinity of Nix, WGPU and GitHub. \ No newline at end of file