From f5c7375bf87ff4b7c1c8dbc81d8daab5ac527649 Mon Sep 17 00:00:00 2001 From: Corey Lowman Date: Sat, 15 Oct 2022 11:29:59 -0400 Subject: [PATCH 1/6] Working commit of new numpy feature --- Cargo.toml | 7 +- examples/06-mnist.rs | 1 + examples/numpy-save-load.rs | 10 ++- src/lib.rs | 1 + src/nn/activations.rs | 6 ++ src/nn/batchnorm2d.rs | 9 +- src/nn/conv.rs | 82 ++++++++++-------- src/nn/dropout.rs | 4 + src/nn/flatten.rs | 2 + src/nn/generalized_residual.rs | 53 +++++++----- src/nn/impl_module_for_tuples.rs | 139 ++++++++++++++++-------------- src/nn/layer_norm.rs | 100 +++++++++++---------- src/nn/linear.rs | 83 ++++++++++-------- src/nn/mod.rs | 8 +- src/nn/pool2d.rs | 6 +- src/nn/pool_global.rs | 8 +- src/nn/repeated.rs | 101 ++++++++++++---------- src/nn/residual.rs | 37 +++++--- src/nn/split_into.rs | 2 + src/nn/transformer/decoder.rs | 6 ++ src/nn/transformer/encoder.rs | 4 + src/nn/transformer/mha.rs | 3 + src/nn/transformer/transformer.rs | 5 ++ 23 files changed, 402 insertions(+), 275 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 293d4ae71..0e8fcde02 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,19 +21,20 @@ keywords = [ # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [package.metadata.docs.rs] -features = ["nightly"] +features = ["nightly", "numpy"] [dependencies] rand = { version = "0.8.5", features = ["std_rng"] } rand_distr = { version = "0.4.3", features = [] } matrixmultiply = { version = "0.3.2", features = [] } -zip = { version = "0.6.2", features = [] } +zip = { version = "0.6.2", features = [], optional = true } cblas-sys = { version = "0.1.4", optional = true } libc = { version = "0.2", optional = true } [features] -default = [] +default = ["numpy"] nightly = [] +numpy = ["dep:zip"] cblas = ["dep:cblas-sys", "dep:libc"] intel-mkl = ["cblas"] diff --git a/examples/06-mnist.rs b/examples/06-mnist.rs index e7e18a0de..5b34b1521 100644 --- a/examples/06-mnist.rs +++ b/examples/06-mnist.rs @@ -107,6 +107,7 @@ fn main() { } // save our model to a .npz file + #[cfg(feature = "numpy")] model .save("mnist-classifier.npz") .expect("failed to save model"); diff --git a/examples/numpy-save-load.rs b/examples/numpy-save-load.rs index 3db82f6cd..bbaeff5d8 100644 --- a/examples/numpy-save-load.rs +++ b/examples/numpy-save-load.rs @@ -1,8 +1,9 @@ //! Demonstrates how to use dfdx::numpy to save and load arrays -use dfdx::numpy as np; - +#[cfg(feature = "numpy")] fn main() { + use dfdx::numpy as np; + np::save("0d-rs.npy", &1.234).expect("Saving failed"); np::save("1d-rs.npy", &[1.0, 2.0, 3.0]).expect("Saving failed"); np::save("2d-rs.npy", &[[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]).expect("Saving failed"); @@ -19,3 +20,8 @@ fn main() { np::load("2d-rs.npy", &mut expected_2d).expect("Loading failed"); assert_eq!(expected_2d, [[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]); } + +#[cfg(not(feature = "numpy"))] +fn main() { + panic!("Use the 'numpy' feature to run this example"); +} diff --git a/src/lib.rs b/src/lib.rs index 62163c7af..c03aeaf7c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -101,6 +101,7 @@ pub mod feature_flags; pub mod gradients; pub mod losses; pub mod nn; +#[cfg(feature = "numpy")] pub mod numpy; pub mod optim; pub mod tensor; diff --git a/src/nn/activations.rs b/src/nn/activations.rs index 136076293..7e6659b6f 100644 --- a/src/nn/activations.rs +++ b/src/nn/activations.rs @@ -19,7 +19,10 @@ macro_rules! activation_impls { fn reset_params(&mut self, _: &mut R) {} } + #[cfg(feature = "numpy")] impl SaveToNpz for $struct_name {} + + #[cfg(feature = "numpy")] impl LoadFromNpz for $struct_name {} impl> Module for $struct_name { @@ -66,7 +69,10 @@ impl ResetParams for Softmax { fn reset_params(&mut self, _: &mut R) {} } +#[cfg(feature = "numpy")] impl SaveToNpz for Softmax {} + +#[cfg(feature = "numpy")] impl LoadFromNpz for Softmax {} impl Module for Softmax diff --git a/src/nn/batchnorm2d.rs b/src/nn/batchnorm2d.rs index 9a3c40266..da57df4a6 100644 --- a/src/nn/batchnorm2d.rs +++ b/src/nn/batchnorm2d.rs @@ -1,9 +1,13 @@ -use super::{npz_fread, npz_fwrite, LoadFromNpz, NpzError, SaveToNpz}; use super::{Module, ModuleMut, ResetParams}; use crate::arrays::{HasArrayData, HasAxes}; use crate::devices::{Cpu, FillElements}; use crate::{gradients::*, tensor::*, tensor_ops::*}; + +#[cfg(feature = "numpy")] +use super::{npz_fread, npz_fwrite, LoadFromNpz, NpzError, SaveToNpz}; +#[cfg(feature = "numpy")] use std::io::{Read, Seek, Write}; +#[cfg(feature = "numpy")] use zip::{result::ZipResult, ZipArchive}; /// Batch normalization for images as described in @@ -191,6 +195,7 @@ impl CanUpdateWithGradients for BatchNorm2D { } } +#[cfg(feature = "numpy")] impl SaveToNpz for BatchNorm2D { fn write(&self, p: &str, w: &mut zip::ZipWriter) -> ZipResult<()> { npz_fwrite(w, format!("{p}scale.npy"), self.scale.data())?; @@ -201,6 +206,7 @@ impl SaveToNpz for BatchNorm2D { } } +#[cfg(feature = "numpy")] impl LoadFromNpz for BatchNorm2D { fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { npz_fread(r, format!("{p}scale.npy"), self.scale.mut_data())?; @@ -323,6 +329,7 @@ mod tests { ); } + #[cfg(feature = "numpy")] #[test] fn test_batchnorm2d_save_load() { let mut rng = StdRng::seed_from_u64(13); diff --git a/src/nn/conv.rs b/src/nn/conv.rs index 7a5f167dc..beac25c5f 100644 --- a/src/nn/conv.rs +++ b/src/nn/conv.rs @@ -2,7 +2,10 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, Tape, UnusedTen use crate::prelude::*; use rand::Rng; use rand_distr::Uniform; + +#[cfg(feature = "numpy")] use std::io::{Read, Seek, Write}; +#[cfg(feature = "numpy")] use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// **Requires Nightly** Performs 2d convolutions on 3d and 4d images. @@ -57,6 +60,7 @@ impl SaveToNpz for Conv2D { @@ -69,6 +73,7 @@ impl LoadFromNpz for Conv2D { @@ -139,8 +144,6 @@ where mod tests { use super::*; use rand::thread_rng; - use std::fs::File; - use tempfile::NamedTempFile; #[test] fn test_forward_3d_sizes() { @@ -187,45 +190,52 @@ mod tests { let _: Tensor3D<1, 8, 8> = <(A, B, C)>::default().forward_mut(Img::zeros()); } - #[test] - fn test_save_conv2d() { - let model: Conv2D<2, 4, 3> = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let mut zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - { - let weight_file = zip - .by_name("weight.npy") - .expect("failed to find weight.npy file"); - assert!(weight_file.size() > 0); - } - { - let bias_file = zip - .by_name("bias.npy") - .expect("failed to find bias.npy file"); - assert!(bias_file.size() > 0); + #[cfg(feature = "numpy")] + mod numpy_tests { + use super::*; + use std::fs::File; + use tempfile::NamedTempFile; + + #[test] + fn test_save_conv2d() { + let model: Conv2D<2, 4, 3> = Default::default(); + let file = NamedTempFile::new().expect("failed to create tempfile"); + model + .save(file.path().to_str().unwrap()) + .expect("failed to save model"); + let f = File::open(file.path()).expect("failed to open resulting file"); + let mut zip = ZipArchive::new(f).expect("failed to create zip archive from file"); + { + let weight_file = zip + .by_name("weight.npy") + .expect("failed to find weight.npy file"); + assert!(weight_file.size() > 0); + } + { + let bias_file = zip + .by_name("bias.npy") + .expect("failed to find bias.npy file"); + assert!(bias_file.size() > 0); + } } - } - #[test] - fn test_load_conv() { - let mut rng = thread_rng(); - let mut saved_model: Conv2D<2, 4, 3> = Default::default(); - saved_model.reset_params(&mut rng); + #[test] + fn test_load_conv() { + let mut rng = thread_rng(); + let mut saved_model: Conv2D<2, 4, 3> = Default::default(); + saved_model.reset_params(&mut rng); - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - let mut loaded_model: Conv2D<2, 4, 3> = Default::default(); - assert!(loaded_model.weight.data() != saved_model.weight.data()); - assert!(loaded_model.bias.data() != saved_model.bias.data()); + let mut loaded_model: Conv2D<2, 4, 3> = Default::default(); + assert!(loaded_model.weight.data() != saved_model.weight.data()); + assert!(loaded_model.bias.data() != saved_model.bias.data()); - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.weight.data(), saved_model.weight.data()); - assert_eq!(loaded_model.bias.data(), saved_model.bias.data()); + assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); + assert_eq!(loaded_model.weight.data(), saved_model.weight.data()); + assert_eq!(loaded_model.bias.data(), saved_model.bias.data()); + } } #[test] diff --git a/src/nn/dropout.rs b/src/nn/dropout.rs index 4d345c05b..31495aa5d 100644 --- a/src/nn/dropout.rs +++ b/src/nn/dropout.rs @@ -68,7 +68,9 @@ impl ResetParams for DropoutOneIn { fn reset_params(&mut self, _: &mut R) {} } +#[cfg(feature = "numpy")] impl SaveToNpz for DropoutOneIn {} +#[cfg(feature = "numpy")] impl LoadFromNpz for DropoutOneIn {} impl> Module for DropoutOneIn { @@ -166,7 +168,9 @@ impl ResetParams for Dropout { fn reset_params(&mut self, _: &mut R) {} } +#[cfg(feature = "numpy")] impl SaveToNpz for Dropout {} +#[cfg(feature = "numpy")] impl LoadFromNpz for Dropout {} impl> Module for Dropout { diff --git a/src/nn/flatten.rs b/src/nn/flatten.rs index bc560b8dc..95953f0db 100644 --- a/src/nn/flatten.rs +++ b/src/nn/flatten.rs @@ -23,7 +23,9 @@ impl CanUpdateWithGradients for FlattenImage { fn update(&mut self, _: &mut G, _: &mut UnusedTensors) {} } +#[cfg(feature = "numpy")] impl SaveToNpz for FlattenImage {} +#[cfg(feature = "numpy")] impl LoadFromNpz for FlattenImage {} impl Module> diff --git a/src/nn/generalized_residual.rs b/src/nn/generalized_residual.rs index 2e7e3b0ce..bea8d617e 100644 --- a/src/nn/generalized_residual.rs +++ b/src/nn/generalized_residual.rs @@ -1,6 +1,8 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, UnusedTensors}; use crate::prelude::*; +#[cfg(feature = "numpy")] use std::io::{Read, Seek, Write}; +#[cfg(feature = "numpy")] use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// A residual connection `R` around `F`: `F(x) + R(x)`, @@ -89,6 +91,7 @@ where } } +#[cfg(feature = "numpy")] impl SaveToNpz for GeneralizedResidual { /// Pass through to `F`/`R`'s [SaveToNpz]. fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { @@ -97,6 +100,7 @@ impl SaveToNpz for GeneralizedResidual { } } +#[cfg(feature = "numpy")] impl LoadFromNpz for GeneralizedResidual { /// Pass through to `F`/`R`'s [LoadFromNpz]. fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { @@ -110,7 +114,6 @@ mod tests { use super::*; use crate::tests::assert_close; use rand::{prelude::StdRng, SeedableRng}; - use tempfile::NamedTempFile; #[test] fn test_reset_generalized_residual() { @@ -149,25 +152,33 @@ mod tests { assert_close(g.ref_gradient(&model.r.bias), &[0.5; 2]); } - #[test] - fn test_save_load_generalized_residual() { - let mut rng = StdRng::seed_from_u64(0); - let mut saved_model: GeneralizedResidual, Linear<5, 3>> = Default::default(); - saved_model.reset_params(&mut rng); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded_model: GeneralizedResidual, Linear<5, 3>> = Default::default(); - assert_ne!(loaded_model.f.weight.data(), saved_model.f.weight.data()); - assert_ne!(loaded_model.f.bias.data(), saved_model.f.bias.data()); - assert_ne!(loaded_model.r.weight.data(), saved_model.r.weight.data()); - assert_ne!(loaded_model.r.bias.data(), saved_model.r.bias.data()); - - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.f.weight.data(), saved_model.f.weight.data()); - assert_eq!(loaded_model.f.bias.data(), saved_model.f.bias.data()); - assert_eq!(loaded_model.r.weight.data(), saved_model.r.weight.data()); - assert_eq!(loaded_model.r.bias.data(), saved_model.r.bias.data()); + #[cfg(feature = "numpy")] + mod numpy_tests { + use super::*; + use tempfile::NamedTempFile; + + #[test] + fn test_save_load_generalized_residual() { + let mut rng = StdRng::seed_from_u64(0); + let mut saved_model: GeneralizedResidual, Linear<5, 3>> = + Default::default(); + saved_model.reset_params(&mut rng); + + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + + let mut loaded_model: GeneralizedResidual, Linear<5, 3>> = + Default::default(); + assert_ne!(loaded_model.f.weight.data(), saved_model.f.weight.data()); + assert_ne!(loaded_model.f.bias.data(), saved_model.f.bias.data()); + assert_ne!(loaded_model.r.weight.data(), saved_model.r.weight.data()); + assert_ne!(loaded_model.r.bias.data(), saved_model.r.bias.data()); + + assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); + assert_eq!(loaded_model.f.weight.data(), saved_model.f.weight.data()); + assert_eq!(loaded_model.f.bias.data(), saved_model.f.bias.data()); + assert_eq!(loaded_model.r.weight.data(), saved_model.r.weight.data()); + assert_eq!(loaded_model.r.bias.data(), saved_model.r.bias.data()); + } } } diff --git a/src/nn/impl_module_for_tuples.rs b/src/nn/impl_module_for_tuples.rs index 3f56b59c4..974daedc6 100644 --- a/src/nn/impl_module_for_tuples.rs +++ b/src/nn/impl_module_for_tuples.rs @@ -1,7 +1,9 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, UnusedTensors}; use crate::prelude::*; use rand::prelude::Rng; +#[cfg(feature = "numpy")] use std::io::{Read, Seek, Write}; +#[cfg(feature = "numpy")] use zip::{result::ZipResult, ZipArchive, ZipWriter}; macro_rules! tuple_impls { @@ -18,6 +20,7 @@ macro_rules! tuple_impls { } } + #[cfg(feature = "numpy")] impl<$($name: SaveToNpz),+> SaveToNpz for ($($name,)+) { /// Calls `SaveToNpz::write(self., ...)` on each part of the tuple. See [SaveToNpz]. /// @@ -30,6 +33,7 @@ macro_rules! tuple_impls { } } + #[cfg(feature = "numpy")] impl<$($name: LoadFromNpz),+> LoadFromNpz for ($($name,)+) { /// Calls `LoadFromNpz::read(self., ...)` on each part of the tuple. See [LoadFromNpz]. /// @@ -111,8 +115,6 @@ mod tests { use crate::nn::tests::SimpleGradients; use crate::unique_id::HasUniqueId; use rand::{prelude::StdRng, SeedableRng}; - use std::fs::File; - use tempfile::NamedTempFile; #[test] fn test_2_tuple() { @@ -158,70 +160,77 @@ mod tests { assert!(model.1.bias.data() != m0.1.bias.data()); } - #[test] - fn test_save_tuple() { - let model: ( - Linear<1, 2>, - ReLU, - Linear<2, 3>, - (Dropout, Linear<1, 2>, Linear<3, 4>), - ) = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - let mut names = zip.file_names().collect::>(); - names.sort_unstable(); - assert_eq!( - &names, - &[ - "0.bias.npy", - "0.weight.npy", - "2.bias.npy", - "2.weight.npy", - "3.1.bias.npy", - "3.1.weight.npy", - "3.2.bias.npy", - "3.2.weight.npy", - ] - ); - } - - #[test] - fn test_load_tuple() { - type Model = ( - Linear<1, 2>, - ReLU, - Linear<2, 3>, - (Dropout, Linear<1, 2>, Linear<3, 4>), - ); - - let mut rng = StdRng::seed_from_u64(0); - let mut saved_model: Model = Default::default(); - saved_model.reset_params(&mut rng); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded_model: Model = Default::default(); - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.0.weight.data(), saved_model.0.weight.data()); - assert_eq!(loaded_model.0.bias.data(), saved_model.0.bias.data()); - assert_eq!(loaded_model.2.weight.data(), saved_model.2.weight.data()); - assert_eq!(loaded_model.2.bias.data(), saved_model.2.bias.data()); - assert_eq!( - loaded_model.3 .1.weight.data(), - saved_model.3 .1.weight.data() - ); - assert_eq!(loaded_model.3 .1.bias.data(), saved_model.3 .1.bias.data()); + #[cfg(feature = "numpy")] + mod numpy_tests { + use super::*; + use std::fs::File; + use tempfile::NamedTempFile; + + #[test] + fn test_save_tuple() { + let model: ( + Linear<1, 2>, + ReLU, + Linear<2, 3>, + (Dropout, Linear<1, 2>, Linear<3, 4>), + ) = Default::default(); + let file = NamedTempFile::new().expect("failed to create tempfile"); + model + .save(file.path().to_str().unwrap()) + .expect("failed to save model"); + let f = File::open(file.path()).expect("failed to open resulting file"); + let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); + let mut names = zip.file_names().collect::>(); + names.sort_unstable(); + assert_eq!( + &names, + &[ + "0.bias.npy", + "0.weight.npy", + "2.bias.npy", + "2.weight.npy", + "3.1.bias.npy", + "3.1.weight.npy", + "3.2.bias.npy", + "3.2.weight.npy", + ] + ); + } - assert_eq!( - loaded_model.3 .2.weight.data(), - saved_model.3 .2.weight.data() - ); - assert_eq!(loaded_model.3 .2.bias.data(), saved_model.3 .2.bias.data()); + #[test] + fn test_load_tuple() { + type Model = ( + Linear<1, 2>, + ReLU, + Linear<2, 3>, + (Dropout, Linear<1, 2>, Linear<3, 4>), + ); + + let mut rng = StdRng::seed_from_u64(0); + let mut saved_model: Model = Default::default(); + saved_model.reset_params(&mut rng); + + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + + let mut loaded_model: Model = Default::default(); + assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); + assert_eq!(loaded_model.0.weight.data(), saved_model.0.weight.data()); + assert_eq!(loaded_model.0.bias.data(), saved_model.0.bias.data()); + assert_eq!(loaded_model.2.weight.data(), saved_model.2.weight.data()); + assert_eq!(loaded_model.2.bias.data(), saved_model.2.bias.data()); + assert_eq!( + loaded_model.3 .1.weight.data(), + saved_model.3 .1.weight.data() + ); + assert_eq!(loaded_model.3 .1.bias.data(), saved_model.3 .1.bias.data()); + + assert_eq!( + loaded_model.3 .2.weight.data(), + saved_model.3 .2.weight.data() + ); + assert_eq!(loaded_model.3 .2.bias.data(), saved_model.3 .2.bias.data()); + } } /// A struct to test the forward method of tuples. This sets the `I`th valuein a 1d tensors of size `N` to 1.0. diff --git a/src/nn/layer_norm.rs b/src/nn/layer_norm.rs index 2efd60fab..6e98b1f4c 100644 --- a/src/nn/layer_norm.rs +++ b/src/nn/layer_norm.rs @@ -2,7 +2,9 @@ use crate::arrays::Axis; use crate::devices::{Cpu, FillElements}; use crate::gradients::{CanUpdateWithGradients, GradientProvider, Tape, UnusedTensors}; use crate::prelude::*; +#[cfg(feature = "numpy")] use std::io::{Read, Seek, Write}; +#[cfg(feature = "numpy")] use zip::{result::ZipResult, ZipArchive}; /// Implements layer normalization as described in [Layer Normalization](https://arxiv.org/abs/1607.06450). @@ -114,6 +116,7 @@ where } } +#[cfg(feature = "numpy")] impl SaveToNpz for LayerNorm1D { /// Saves [Self::gamma] to `{pre}gamma.npy` and [Self::beta] to `{pre}beta.npy` /// using [npz_fwrite()]. @@ -124,6 +127,7 @@ impl SaveToNpz for LayerNorm1D { } } +#[cfg(feature = "numpy")] impl LoadFromNpz for LayerNorm1D { /// Reads [Self::gamma] from `{p}gamma.npy` and [Self::beta] from `{p}beta.npy` /// using [npz_fread()]. @@ -141,8 +145,6 @@ mod tests { use crate::{nn::tests::SimpleGradients, tests::assert_close}; use rand::{prelude::StdRng, SeedableRng}; use rand_distr::Standard; - use std::fs::File; - use tempfile::NamedTempFile; #[test] fn test_layer_norm_reset() { @@ -202,56 +204,62 @@ mod tests { assert_close(g.ref_gradient(&m.beta), &[0.2; 5]); } - #[test] - fn test_save_layer_norm() { - let model: LayerNorm1D<13> = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - let mut names = zip.file_names().collect::>(); - names.sort_unstable(); - assert_eq!(&names, &["beta.npy", "gamma.npy",]); - } + #[cfg(feature = "numpy")] + mod numpy_tests { + use super::*; + use std::fs::File; + use tempfile::NamedTempFile; + + #[test] + fn test_save_layer_norm() { + let model: LayerNorm1D<13> = Default::default(); + let file = NamedTempFile::new().expect("failed to create tempfile"); + model + .save(file.path().to_str().unwrap()) + .expect("failed to save model"); + let f = File::open(file.path()).expect("failed to open resulting file"); + let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); + let mut names = zip.file_names().collect::>(); + names.sort_unstable(); + assert_eq!(&names, &["beta.npy", "gamma.npy",]); + } - #[test] - fn test_save_layer_norm_tuple() { - let model: (LayerNorm1D<5>, LayerNorm1D<13>) = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - let mut names = zip.file_names().collect::>(); - names.sort_unstable(); - assert_eq!( - &names, - &["0.beta.npy", "0.gamma.npy", "1.beta.npy", "1.gamma.npy"] - ); - } + #[test] + fn test_save_layer_norm_tuple() { + let model: (LayerNorm1D<5>, LayerNorm1D<13>) = Default::default(); + let file = NamedTempFile::new().expect("failed to create tempfile"); + model + .save(file.path().to_str().unwrap()) + .expect("failed to save model"); + let f = File::open(file.path()).expect("failed to open resulting file"); + let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); + let mut names = zip.file_names().collect::>(); + names.sort_unstable(); + assert_eq!( + &names, + &["0.beta.npy", "0.gamma.npy", "1.beta.npy", "1.gamma.npy"] + ); + } - #[test] - fn test_load_layer_norm() { - let mut rng = StdRng::seed_from_u64(0); - let mut saved_model: LayerNorm1D<13> = Default::default(); - saved_model.gamma.randomize(&mut rng, &Standard); - saved_model.beta.randomize(&mut rng, &Standard); + #[test] + fn test_load_layer_norm() { + let mut rng = StdRng::seed_from_u64(0); + let mut saved_model: LayerNorm1D<13> = Default::default(); + saved_model.gamma.randomize(&mut rng, &Standard); + saved_model.beta.randomize(&mut rng, &Standard); - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - let mut loaded_model: LayerNorm1D<13> = Default::default(); - assert!(loaded_model.gamma.data() != saved_model.gamma.data()); - assert!(loaded_model.beta.data() != saved_model.beta.data()); + let mut loaded_model: LayerNorm1D<13> = Default::default(); + assert!(loaded_model.gamma.data() != saved_model.gamma.data()); + assert!(loaded_model.beta.data() != saved_model.beta.data()); - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.gamma.data(), saved_model.gamma.data()); - assert_eq!(loaded_model.beta.data(), saved_model.beta.data()); + assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); + assert_eq!(loaded_model.gamma.data(), saved_model.gamma.data()); + assert_eq!(loaded_model.beta.data(), saved_model.beta.data()); + } } - #[test] fn test_layer_norm_missing_gradients() { let mut model: LayerNorm1D<5> = Default::default(); diff --git a/src/nn/linear.rs b/src/nn/linear.rs index ad115da6c..c13ea544d 100644 --- a/src/nn/linear.rs +++ b/src/nn/linear.rs @@ -2,7 +2,9 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, Tape, UnusedTen use crate::prelude::*; use rand::Rng; use rand_distr::Uniform; +#[cfg(feature = "numpy")] use std::io::{Read, Seek, Write}; +#[cfg(feature = "numpy")] use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// A linear transformation of the form `weight * x + bias`, where `weight` is a matrix, `x` is a vector or matrix, @@ -52,6 +54,7 @@ impl ResetParams for Linear { } } +#[cfg(feature = "numpy")] impl SaveToNpz for Linear { /// Saves [Self::weight] to `{pre}weight.npy` and [Self::bias] to `{pre}bias.npy` /// using [npz_fwrite()]. @@ -65,6 +68,7 @@ impl SaveToNpz for Linear { } } +#[cfg(feature = "numpy")] impl LoadFromNpz for Linear { /// Reads [Self::weight] from `{pre}weight.npy` and [Self::bias] from `{pre}bias.npy` /// using [npz_fread()]. @@ -126,9 +130,6 @@ mod tests { use super::*; use crate::unique_id::HasUniqueId; use crate::{nn::tests::SimpleGradients, tests::assert_close}; - use rand::{prelude::StdRng, SeedableRng}; - use std::fs::File; - use tempfile::NamedTempFile; const W: [[f32; 5]; 2] = [ [-0.3458893, -0.30371523, -0.3712057, 0.14303583, -0.0268966], @@ -257,45 +258,53 @@ mod tests { ); } - #[test] - fn test_save_linear() { - let model: Linear<5, 3> = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let mut zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - { - let weight_file = zip - .by_name("weight.npy") - .expect("failed to find weight.npy file"); - assert!(weight_file.size() > 0); - } - { - let bias_file = zip - .by_name("bias.npy") - .expect("failed to find bias.npy file"); - assert!(bias_file.size() > 0); + #[cfg(feature = "numpy")] + mod numpy_tests { + use super::*; + use rand::{prelude::StdRng, SeedableRng}; + use std::fs::File; + use tempfile::NamedTempFile; + + #[test] + fn test_save_linear() { + let model: Linear<5, 3> = Default::default(); + let file = NamedTempFile::new().expect("failed to create tempfile"); + model + .save(file.path().to_str().unwrap()) + .expect("failed to save model"); + let f = File::open(file.path()).expect("failed to open resulting file"); + let mut zip = ZipArchive::new(f).expect("failed to create zip archive from file"); + { + let weight_file = zip + .by_name("weight.npy") + .expect("failed to find weight.npy file"); + assert!(weight_file.size() > 0); + } + { + let bias_file = zip + .by_name("bias.npy") + .expect("failed to find bias.npy file"); + assert!(bias_file.size() > 0); + } } - } - #[test] - fn test_load_linear() { - let mut rng = StdRng::seed_from_u64(0); - let mut saved_model: Linear<5, 3> = Default::default(); - saved_model.reset_params(&mut rng); + #[test] + fn test_load_linear() { + let mut rng = StdRng::seed_from_u64(0); + let mut saved_model: Linear<5, 3> = Default::default(); + saved_model.reset_params(&mut rng); - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - let mut loaded_model: Linear<5, 3> = Default::default(); - assert!(loaded_model.weight.data() != saved_model.weight.data()); - assert!(loaded_model.bias.data() != saved_model.bias.data()); + let mut loaded_model: Linear<5, 3> = Default::default(); + assert!(loaded_model.weight.data() != saved_model.weight.data()); + assert!(loaded_model.bias.data() != saved_model.bias.data()); - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.weight.data(), saved_model.weight.data()); - assert_eq!(loaded_model.bias.data(), saved_model.bias.data()); + assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); + assert_eq!(loaded_model.weight.data(), saved_model.weight.data()); + assert_eq!(loaded_model.bias.data(), saved_model.bias.data()); + } } #[test] diff --git a/src/nn/mod.rs b/src/nn/mod.rs index 4030ce968..d4bc1b43b 100644 --- a/src/nn/mod.rs +++ b/src/nn/mod.rs @@ -84,7 +84,6 @@ mod impl_module_for_tuples; mod layer_norm; mod linear; mod module; -mod npz; mod pool_global; mod repeated; mod residual; @@ -98,12 +97,17 @@ pub use impl_module_for_tuples::*; pub use layer_norm::*; pub use linear::*; pub use module::*; -pub use npz::*; pub use pool_global::*; pub use repeated::*; pub use residual::*; pub use split_into::*; +#[cfg(feature = "numpy")] +mod npz; + +#[cfg(feature = "numpy")] +pub use npz::*; + #[cfg(feature = "nightly")] mod transformer; diff --git a/src/nn/pool2d.rs b/src/nn/pool2d.rs index d2daceaf3..9fda878d6 100644 --- a/src/nn/pool2d.rs +++ b/src/nn/pool2d.rs @@ -1,4 +1,6 @@ -use super::{LoadFromNpz, Module, ModuleMut, ResetParams, SaveToNpz}; +#[cfg(feature = "numpy")] +use super::{LoadFromNpz, SaveToNpz}; +use super::{Module, ModuleMut, ResetParams}; use crate::gradients::*; use crate::tensor::*; use rand::Rng; @@ -45,7 +47,9 @@ macro_rules! impl_pools { fn reset_params(&mut self, _: &mut R) {} } + #[cfg(feature = "numpy")] impl SaveToNpz for $PoolTy {} + #[cfg(feature = "numpy")] impl LoadFromNpz for $PoolTy {} impl< diff --git a/src/nn/pool_global.rs b/src/nn/pool_global.rs index d8290a1d4..233139966 100644 --- a/src/nn/pool_global.rs +++ b/src/nn/pool_global.rs @@ -1,4 +1,6 @@ -use super::{LoadFromNpz, Module, ModuleMut, ResetParams, SaveToNpz}; +#[cfg(feature = "numpy")] +use super::{LoadFromNpz, SaveToNpz}; +use super::{Module, ModuleMut, ResetParams}; use crate::gradients::*; use crate::tensor::*; @@ -64,7 +66,11 @@ macro_rules! impl_pools { impl CanUpdateWithGradients for $PoolTy { fn update(&mut self, _: &mut G, _: &mut UnusedTensors) {} } + + #[cfg(feature = "numpy")] impl SaveToNpz for $PoolTy {} + + #[cfg(feature = "numpy")] impl LoadFromNpz for $PoolTy {} impl Module> for $PoolTy { diff --git a/src/nn/repeated.rs b/src/nn/repeated.rs index 04516c5bf..d75894718 100644 --- a/src/nn/repeated.rs +++ b/src/nn/repeated.rs @@ -1,6 +1,8 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, UnusedTensors}; use crate::prelude::*; +#[cfg(feature = "numpy")] use std::io::{Read, Seek, Write}; +#[cfg(feature = "numpy")] use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// Repeats `T` `N` times. This requires that `T`'s input is the same as it's output. @@ -54,6 +56,7 @@ impl CanUpdateWithGradients for Repea } } +#[cfg(feature = "numpy")] impl SaveToNpz for Repeated { /// Calls `SaveToNpz::write(self.modules[i], ...)` on each sub module. See [SaveToNpz]. /// @@ -68,6 +71,7 @@ impl SaveToNpz for Repeated { } } +#[cfg(feature = "numpy")] impl LoadFromNpz for Repeated { /// Calls `LoadFromNpz::read(self.modules[i], ...)` on each sub module. See [LoadFromNpz]. /// @@ -113,8 +117,6 @@ mod tests { use crate::nn::tests::SimpleGradients; use crate::unique_id::HasUniqueId; use rand::{prelude::StdRng, SeedableRng}; - use std::fs::File; - use tempfile::NamedTempFile; #[test] fn test_default_and_reset() { @@ -154,54 +156,61 @@ mod tests { assert_eq!(x.data(), m.forward_mut(Tensor1D::zeros()).data()); } - #[test] - fn test_save_repeated() { - let model: Repeated, 4> = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - let mut names = zip.file_names().collect::>(); - names.sort_unstable(); - assert_eq!( - &names, - &[ - "0.bias.npy", - "0.weight.npy", - "1.bias.npy", - "1.weight.npy", - "2.bias.npy", - "2.weight.npy", - "3.bias.npy", - "3.weight.npy", - ] - ); - } + #[cfg(feature = "numpy")] + mod numpy_tests { + use super::*; + use std::fs::File; + use tempfile::NamedTempFile; - #[test] - fn test_load_repeated() { - type Model = Repeated, 4>; + #[test] + fn test_save_repeated() { + let model: Repeated, 4> = Default::default(); + let file = NamedTempFile::new().expect("failed to create tempfile"); + model + .save(file.path().to_str().unwrap()) + .expect("failed to save model"); + let f = File::open(file.path()).expect("failed to open resulting file"); + let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); + let mut names = zip.file_names().collect::>(); + names.sort_unstable(); + assert_eq!( + &names, + &[ + "0.bias.npy", + "0.weight.npy", + "1.bias.npy", + "1.weight.npy", + "2.bias.npy", + "2.weight.npy", + "3.bias.npy", + "3.weight.npy", + ] + ); + } - let mut rng = StdRng::seed_from_u64(0); - let mut saved_model: Model = Default::default(); - saved_model.reset_params(&mut rng); + #[test] + fn test_load_repeated() { + type Model = Repeated, 4>; - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + let mut rng = StdRng::seed_from_u64(0); + let mut saved_model: Model = Default::default(); + saved_model.reset_params(&mut rng); - let mut loaded_model: Model = Default::default(); - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - for i in 0..4 { - assert_eq!( - loaded_model.modules[i].weight.data(), - saved_model.modules[i].weight.data() - ); - assert_eq!( - loaded_model.modules[i].bias.data(), - saved_model.modules[i].bias.data() - ); + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + + let mut loaded_model: Model = Default::default(); + assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); + for i in 0..4 { + assert_eq!( + loaded_model.modules[i].weight.data(), + saved_model.modules[i].weight.data() + ); + assert_eq!( + loaded_model.modules[i].bias.data(), + saved_model.modules[i].bias.data() + ); + } } } diff --git a/src/nn/residual.rs b/src/nn/residual.rs index 8a7f869d2..cb3840576 100644 --- a/src/nn/residual.rs +++ b/src/nn/residual.rs @@ -1,6 +1,8 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, UnusedTensors}; use crate::prelude::*; +#[cfg(feature = "numpy")] use std::io::{Read, Seek, Write}; +#[cfg(feature = "numpy")] use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// A residual connection around `F`: `F(x) + x`, @@ -50,6 +52,7 @@ impl, F: ModuleMut> ModuleMut for Resid } } +#[cfg(feature = "numpy")] impl SaveToNpz for Residual { /// Pass through to `F`'s [SaveToNpz]. fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { @@ -57,6 +60,7 @@ impl SaveToNpz for Residual { } } +#[cfg(feature = "numpy")] impl LoadFromNpz for Residual { /// Pass through to `F`'s [LoadFromNpz]. fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { @@ -69,7 +73,6 @@ mod tests { use super::*; use crate::tests::assert_close; use rand::{prelude::StdRng, SeedableRng}; - use tempfile::NamedTempFile; #[test] fn test_residual_reset() { @@ -102,21 +105,27 @@ mod tests { assert_close(g.ref_gradient(&x), &[[0.18806472, 0.21419683]; 4]); } - #[test] - fn test_save_load_residual() { - let mut rng = StdRng::seed_from_u64(0); - let mut saved_model: Residual> = Default::default(); - saved_model.reset_params(&mut rng); + #[cfg(feature = "numpy")] + mod numpy_tests { + use super::*; + use tempfile::NamedTempFile; + + #[test] + fn test_save_load_residual() { + let mut rng = StdRng::seed_from_u64(0); + let mut saved_model: Residual> = Default::default(); + saved_model.reset_params(&mut rng); - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - let mut loaded_model: Residual> = Default::default(); - assert_ne!(loaded_model.0.weight.data(), saved_model.0.weight.data()); - assert_ne!(loaded_model.0.bias.data(), saved_model.0.bias.data()); + let mut loaded_model: Residual> = Default::default(); + assert_ne!(loaded_model.0.weight.data(), saved_model.0.weight.data()); + assert_ne!(loaded_model.0.bias.data(), saved_model.0.bias.data()); - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.0.weight.data(), saved_model.0.weight.data()); - assert_eq!(loaded_model.0.bias.data(), saved_model.0.bias.data()); + assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); + assert_eq!(loaded_model.0.weight.data(), saved_model.0.weight.data()); + assert_eq!(loaded_model.0.bias.data(), saved_model.0.bias.data()); + } } } diff --git a/src/nn/split_into.rs b/src/nn/split_into.rs index 4f033b15b..a542ee210 100644 --- a/src/nn/split_into.rs +++ b/src/nn/split_into.rs @@ -32,6 +32,7 @@ impl ResetParams for SplitInto { } } +#[cfg(feature = "numpy")] impl SaveToNpz for SplitInto { fn write(&self, p: &str, w: &mut zip::ZipWriter) -> zip::result::ZipResult<()> where @@ -41,6 +42,7 @@ impl SaveToNpz for SplitInto { } } +#[cfg(feature = "numpy")] impl LoadFromNpz for SplitInto { fn read(&mut self, p: &str, r: &mut zip::ZipArchive) -> Result<(), NpzError> where diff --git a/src/nn/transformer/decoder.rs b/src/nn/transformer/decoder.rs index 396d93551..0373712c0 100644 --- a/src/nn/transformer/decoder.rs +++ b/src/nn/transformer/decoder.rs @@ -1,7 +1,9 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, UnusedTensors}; use crate::prelude::*; use rand::Rng; +#[cfg(feature = "numpy")] use std::io::{Read, Seek, Write}; +#[cfg(feature = "numpy")] use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// **Requires Nightly** A transformer decoder. @@ -65,6 +67,7 @@ where } } +#[cfg(feature = "numpy")] impl SaveToNpz for TransformerDecoder { @@ -73,6 +76,7 @@ impl SaveToNpz } } +#[cfg(feature = "numpy")] impl LoadFromNpz for TransformerDecoder { @@ -169,6 +173,7 @@ where } } +#[cfg(feature = "numpy")] impl SaveToNpz for TransformerDecoderBlock { @@ -184,6 +189,7 @@ impl SaveToNpz } } +#[cfg(feature = "numpy")] impl LoadFromNpz for TransformerDecoderBlock { diff --git a/src/nn/transformer/encoder.rs b/src/nn/transformer/encoder.rs index 81128cfab..c8273509f 100644 --- a/src/nn/transformer/encoder.rs +++ b/src/nn/transformer/encoder.rs @@ -1,6 +1,8 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, UnusedTensors}; use crate::prelude::*; +#[cfg(feature = "numpy")] use std::io::{Read, Seek, Write}; +#[cfg(feature = "numpy")] use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// **Requires Nightly** A transformer encoder. @@ -105,6 +107,7 @@ where } } +#[cfg(feature = "numpy")] impl SaveToNpz for TransformerEncoderBlock { @@ -118,6 +121,7 @@ impl SaveToNpz } } +#[cfg(feature = "numpy")] impl LoadFromNpz for TransformerEncoderBlock { diff --git a/src/nn/transformer/mha.rs b/src/nn/transformer/mha.rs index a90aedc97..1c55a4538 100644 --- a/src/nn/transformer/mha.rs +++ b/src/nn/transformer/mha.rs @@ -53,6 +53,7 @@ impl CanUpdateWi } } +#[cfg(feature = "numpy")] impl SaveToNpz for MultiHeadAttention { @@ -68,6 +69,7 @@ impl SaveToNpz } } +#[cfg(feature = "numpy")] impl LoadFromNpz for MultiHeadAttention { @@ -326,6 +328,7 @@ mod tests { assert!(unused.is_empty()); } + #[cfg(feature = "numpy")] #[test] fn test_save_and_load() { let mut rng = thread_rng(); diff --git a/src/nn/transformer/transformer.rs b/src/nn/transformer/transformer.rs index 12bdddd16..74206d46d 100644 --- a/src/nn/transformer/transformer.rs +++ b/src/nn/transformer/transformer.rs @@ -1,6 +1,8 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, UnusedTensors}; use crate::prelude::*; +#[cfg(feature = "numpy")] use std::io::{Read, Seek, Write}; +#[cfg(feature = "numpy")] use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// **Requires Nightly** Transformer architecture as described in @@ -86,6 +88,7 @@ where } } +#[cfg(feature = "numpy")] impl SaveToNpz for Transformer { @@ -96,6 +99,7 @@ impl LoadFromNpz for Transformer { @@ -148,6 +152,7 @@ mod tests { assert!(unused.is_empty()); } + #[cfg(feature = "numpy")] #[test] fn test_save_load() { let mut rng = thread_rng(); From ee5d63ae2e1c427c0ac3486fbb05a6509426192e Mon Sep 17 00:00:00 2001 From: Corey Lowman Date: Sat, 15 Oct 2022 11:37:35 -0400 Subject: [PATCH 2/6] Adding numpy to feature flags doc --- src/feature_flags.rs | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/feature_flags.rs b/src/feature_flags.rs index a2ed34af6..1af377e91 100644 --- a/src/feature_flags.rs +++ b/src/feature_flags.rs @@ -24,6 +24,15 @@ //! //! `build.rs` will fail helpfully if you don't have the correct path/environment variables. //! +//! # "numpy" +//! +//! Enables saving and loading arrays to .npy files, and saving and loading nn to .npz files. +//! +//! Example: +//! ```toml +//! dfdx = { version = "...", features = ["numpy"] } +//! ``` +//! //! # "nightly" //! //! Enables using all features that currently require the nightly rust compiler. From acc04416988da659605ff9677d547a54a2ace4f8 Mon Sep 17 00:00:00 2001 From: Corey Lowman Date: Sat, 15 Oct 2022 12:05:56 -0400 Subject: [PATCH 3/6] Moving all npz impls to nn/npz_impls.rs --- src/nn/activations.rs | 12 - src/nn/batchnorm2d.rs | 65 +-- src/nn/conv.rs | 79 --- src/nn/dropout.rs | 10 - src/nn/flatten.rs | 5 - src/nn/generalized_residual.rs | 52 -- src/nn/impl_module_for_tuples.rs | 103 ---- src/nn/layer_norm.rs | 82 ---- src/nn/linear.rs | 81 ---- src/nn/mod.rs | 3 + src/nn/npz_impls.rs | 774 ++++++++++++++++++++++++++++++ src/nn/pool2d.rs | 7 - src/nn/pool_global.rs | 8 - src/nn/repeated.rs | 95 ---- src/nn/residual.rs | 44 -- src/nn/split_into.rs | 20 - src/nn/transformer/decoder.rs | 54 --- src/nn/transformer/encoder.rs | 40 +- src/nn/transformer/mha.rs | 56 --- src/nn/transformer/transformer.rs | 56 +-- 20 files changed, 785 insertions(+), 861 deletions(-) create mode 100644 src/nn/npz_impls.rs diff --git a/src/nn/activations.rs b/src/nn/activations.rs index 7e6659b6f..2a34f08c7 100644 --- a/src/nn/activations.rs +++ b/src/nn/activations.rs @@ -19,12 +19,6 @@ macro_rules! activation_impls { fn reset_params(&mut self, _: &mut R) {} } - #[cfg(feature = "numpy")] - impl SaveToNpz for $struct_name {} - - #[cfg(feature = "numpy")] - impl LoadFromNpz for $struct_name {} - impl> Module for $struct_name { type Output = T; fn forward(&self, input: T) -> Self::Output { @@ -69,12 +63,6 @@ impl ResetParams for Softmax { fn reset_params(&mut self, _: &mut R) {} } -#[cfg(feature = "numpy")] -impl SaveToNpz for Softmax {} - -#[cfg(feature = "numpy")] -impl LoadFromNpz for Softmax {} - impl Module for Softmax where T: Reduce<<::Array as HasLastAxis>::LastAxis>, diff --git a/src/nn/batchnorm2d.rs b/src/nn/batchnorm2d.rs index da57df4a6..1940e92cd 100644 --- a/src/nn/batchnorm2d.rs +++ b/src/nn/batchnorm2d.rs @@ -3,13 +3,6 @@ use crate::arrays::{HasArrayData, HasAxes}; use crate::devices::{Cpu, FillElements}; use crate::{gradients::*, tensor::*, tensor_ops::*}; -#[cfg(feature = "numpy")] -use super::{npz_fread, npz_fwrite, LoadFromNpz, NpzError, SaveToNpz}; -#[cfg(feature = "numpy")] -use std::io::{Read, Seek, Write}; -#[cfg(feature = "numpy")] -use zip::{result::ZipResult, ZipArchive}; - /// Batch normalization for images as described in /// [Batch Normalization: Accelerating Deep Network Training /// by Reducing Internal Covariate Shift](https://arxiv.org/abs/1502.03167) @@ -195,36 +188,11 @@ impl CanUpdateWithGradients for BatchNorm2D { } } -#[cfg(feature = "numpy")] -impl SaveToNpz for BatchNorm2D { - fn write(&self, p: &str, w: &mut zip::ZipWriter) -> ZipResult<()> { - npz_fwrite(w, format!("{p}scale.npy"), self.scale.data())?; - npz_fwrite(w, format!("{p}bias.npy"), self.bias.data())?; - npz_fwrite(w, format!("{p}running_mean.npy"), self.running_mean.data())?; - npz_fwrite(w, format!("{p}running_var.npy"), self.running_var.data())?; - Ok(()) - } -} - -#[cfg(feature = "numpy")] -impl LoadFromNpz for BatchNorm2D { - fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - npz_fread(r, format!("{p}scale.npy"), self.scale.mut_data())?; - npz_fread(r, format!("{p}bias.npy"), self.bias.mut_data())?; - let mean = self.running_mean.mut_data(); - npz_fread(r, format!("{p}running_mean.npy"), mean)?; - let var = self.running_var.mut_data(); - npz_fread(r, format!("{p}running_var.npy"), var)?; - Ok(()) - } -} - #[cfg(test)] mod tests { use super::*; - use crate::{nn::tests::SimpleGradients, tests::assert_close}; + use crate::tests::assert_close; use rand::{rngs::StdRng, SeedableRng}; - use tempfile::NamedTempFile; #[test] fn test_batchnorm2d_3d_forward_mut() { @@ -328,35 +296,4 @@ mod tests { ], ); } - - #[cfg(feature = "numpy")] - #[test] - fn test_batchnorm2d_save_load() { - let mut rng = StdRng::seed_from_u64(13); - let mut bn: BatchNorm2D<3> = Default::default(); - - assert_eq!(bn.running_mean.data(), &[0.0; 3]); - assert_eq!(bn.running_var.data(), &[1.0; 3]); - assert_eq!(bn.scale.data(), &[1.0; 3]); - assert_eq!(bn.bias.data(), &[0.0; 3]); - - let x1: Tensor3D<3, 4, 5> = TensorCreator::randn(&mut rng); - let g = backward(bn.forward_mut(x1.trace()).exp().mean()); - bn.update(&mut SimpleGradients(g), &mut Default::default()); - - assert_ne!(bn.running_mean.data(), &[0.0; 3]); - assert_ne!(bn.running_var.data(), &[1.0; 3]); - assert_ne!(bn.scale.data(), &[1.0; 3]); - assert_ne!(bn.bias.data(), &[0.0; 3]); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(bn.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded: BatchNorm2D<3> = Default::default(); - assert!(loaded.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded.scale.data(), bn.scale.data()); - assert_eq!(loaded.bias.data(), bn.bias.data()); - assert_eq!(loaded.running_mean.data(), bn.running_mean.data()); - assert_eq!(loaded.running_var.data(), bn.running_var.data()); - } } diff --git a/src/nn/conv.rs b/src/nn/conv.rs index beac25c5f..0a458567e 100644 --- a/src/nn/conv.rs +++ b/src/nn/conv.rs @@ -3,11 +3,6 @@ use crate::prelude::*; use rand::Rng; use rand_distr::Uniform; -#[cfg(feature = "numpy")] -use std::io::{Read, Seek, Write}; -#[cfg(feature = "numpy")] -use zip::{result::ZipResult, ZipArchive, ZipWriter}; - /// **Requires Nightly** Performs 2d convolutions on 3d and 4d images. /// /// **Pytorch Equivalent**: `torch.nn.Conv2d` @@ -60,32 +55,6 @@ impl SaveToNpz - for Conv2D -{ - /// Saves [Self::weight] to `{pre}weight.npy` and [Self::bias] to `{pre}bias.npy` - /// using [npz_fwrite()]. - fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { - npz_fwrite(w, format!("{pre}weight.npy"), self.weight.data())?; - npz_fwrite(w, format!("{pre}bias.npy"), self.bias.data())?; - Ok(()) - } -} - -#[cfg(feature = "numpy")] -impl LoadFromNpz - for Conv2D -{ - /// Reads [Self::weight] from `{pre}weight.npy` and [Self::bias] from `{pre}bias.npy` - /// using [npz_fread()]. - fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - npz_fread(r, format!("{pre}weight.npy"), self.weight.mut_data())?; - npz_fread(r, format!("{pre}bias.npy"), self.bias.mut_data())?; - Ok(()) - } -} - impl< T: Tape, const I: usize, @@ -190,54 +159,6 @@ mod tests { let _: Tensor3D<1, 8, 8> = <(A, B, C)>::default().forward_mut(Img::zeros()); } - #[cfg(feature = "numpy")] - mod numpy_tests { - use super::*; - use std::fs::File; - use tempfile::NamedTempFile; - - #[test] - fn test_save_conv2d() { - let model: Conv2D<2, 4, 3> = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let mut zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - { - let weight_file = zip - .by_name("weight.npy") - .expect("failed to find weight.npy file"); - assert!(weight_file.size() > 0); - } - { - let bias_file = zip - .by_name("bias.npy") - .expect("failed to find bias.npy file"); - assert!(bias_file.size() > 0); - } - } - - #[test] - fn test_load_conv() { - let mut rng = thread_rng(); - let mut saved_model: Conv2D<2, 4, 3> = Default::default(); - saved_model.reset_params(&mut rng); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded_model: Conv2D<2, 4, 3> = Default::default(); - assert!(loaded_model.weight.data() != saved_model.weight.data()); - assert!(loaded_model.bias.data() != saved_model.bias.data()); - - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.weight.data(), saved_model.weight.data()); - assert_eq!(loaded_model.bias.data(), saved_model.bias.data()); - } - } - #[test] fn test_conv_with_optimizer() { let mut rng = thread_rng(); diff --git a/src/nn/dropout.rs b/src/nn/dropout.rs index 31495aa5d..791e96fcf 100644 --- a/src/nn/dropout.rs +++ b/src/nn/dropout.rs @@ -68,11 +68,6 @@ impl ResetParams for DropoutOneIn { fn reset_params(&mut self, _: &mut R) {} } -#[cfg(feature = "numpy")] -impl SaveToNpz for DropoutOneIn {} -#[cfg(feature = "numpy")] -impl LoadFromNpz for DropoutOneIn {} - impl> Module for DropoutOneIn { type Output = T; /// Does nothing @@ -168,11 +163,6 @@ impl ResetParams for Dropout { fn reset_params(&mut self, _: &mut R) {} } -#[cfg(feature = "numpy")] -impl SaveToNpz for Dropout {} -#[cfg(feature = "numpy")] -impl LoadFromNpz for Dropout {} - impl> Module for Dropout { type Output = T; /// Does nothing. diff --git a/src/nn/flatten.rs b/src/nn/flatten.rs index 95953f0db..ce863f969 100644 --- a/src/nn/flatten.rs +++ b/src/nn/flatten.rs @@ -23,11 +23,6 @@ impl CanUpdateWithGradients for FlattenImage { fn update(&mut self, _: &mut G, _: &mut UnusedTensors) {} } -#[cfg(feature = "numpy")] -impl SaveToNpz for FlattenImage {} -#[cfg(feature = "numpy")] -impl LoadFromNpz for FlattenImage {} - impl Module> for FlattenImage where diff --git a/src/nn/generalized_residual.rs b/src/nn/generalized_residual.rs index bea8d617e..efa722a6b 100644 --- a/src/nn/generalized_residual.rs +++ b/src/nn/generalized_residual.rs @@ -1,9 +1,5 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, UnusedTensors}; use crate::prelude::*; -#[cfg(feature = "numpy")] -use std::io::{Read, Seek, Write}; -#[cfg(feature = "numpy")] -use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// A residual connection `R` around `F`: `F(x) + R(x)`, /// as introduced in [Deep Residual Learning for Image Recognition](https://arxiv.org/abs/1512.03385). @@ -91,24 +87,6 @@ where } } -#[cfg(feature = "numpy")] -impl SaveToNpz for GeneralizedResidual { - /// Pass through to `F`/`R`'s [SaveToNpz]. - fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { - self.f.write(&format!("{p}.f"), w)?; - self.r.write(&format!("{p}.r"), w) - } -} - -#[cfg(feature = "numpy")] -impl LoadFromNpz for GeneralizedResidual { - /// Pass through to `F`/`R`'s [LoadFromNpz]. - fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - self.f.read(&format!("{p}.f"), r)?; - self.r.read(&format!("{p}.r"), r) - } -} - #[cfg(test)] mod tests { use super::*; @@ -151,34 +129,4 @@ mod tests { assert_close(g.ref_gradient(&model.r.weight), &[[-0.025407, 0.155879]; 2]); assert_close(g.ref_gradient(&model.r.bias), &[0.5; 2]); } - - #[cfg(feature = "numpy")] - mod numpy_tests { - use super::*; - use tempfile::NamedTempFile; - - #[test] - fn test_save_load_generalized_residual() { - let mut rng = StdRng::seed_from_u64(0); - let mut saved_model: GeneralizedResidual, Linear<5, 3>> = - Default::default(); - saved_model.reset_params(&mut rng); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded_model: GeneralizedResidual, Linear<5, 3>> = - Default::default(); - assert_ne!(loaded_model.f.weight.data(), saved_model.f.weight.data()); - assert_ne!(loaded_model.f.bias.data(), saved_model.f.bias.data()); - assert_ne!(loaded_model.r.weight.data(), saved_model.r.weight.data()); - assert_ne!(loaded_model.r.bias.data(), saved_model.r.bias.data()); - - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.f.weight.data(), saved_model.f.weight.data()); - assert_eq!(loaded_model.f.bias.data(), saved_model.f.bias.data()); - assert_eq!(loaded_model.r.weight.data(), saved_model.r.weight.data()); - assert_eq!(loaded_model.r.bias.data(), saved_model.r.bias.data()); - } - } } diff --git a/src/nn/impl_module_for_tuples.rs b/src/nn/impl_module_for_tuples.rs index 974daedc6..3640a2b1b 100644 --- a/src/nn/impl_module_for_tuples.rs +++ b/src/nn/impl_module_for_tuples.rs @@ -1,10 +1,6 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, UnusedTensors}; use crate::prelude::*; use rand::prelude::Rng; -#[cfg(feature = "numpy")] -use std::io::{Read, Seek, Write}; -#[cfg(feature = "numpy")] -use zip::{result::ZipResult, ZipArchive, ZipWriter}; macro_rules! tuple_impls { ([$($name:ident),+] [$($idx:tt),+], $last:ident, [$($rev_tail:ident),+]) => { @@ -20,32 +16,6 @@ macro_rules! tuple_impls { } } - #[cfg(feature = "numpy")] - impl<$($name: SaveToNpz),+> SaveToNpz for ($($name,)+) { - /// Calls `SaveToNpz::write(self., ...)` on each part of the tuple. See [SaveToNpz]. - /// - /// E.g. for a two tuple (A, B) with `base == ""`, this will call: - /// 1. `self.0.write("0.", w)` - /// 2. `self.1.write("1.", w)` - fn write(&self, base: &str, w: &mut ZipWriter) -> ZipResult<()> { - $(self.$idx.write(&format!("{}{}.", base, $idx), w)?;)+ - Ok(()) - } - } - - #[cfg(feature = "numpy")] - impl<$($name: LoadFromNpz),+> LoadFromNpz for ($($name,)+) { - /// Calls `LoadFromNpz::read(self., ...)` on each part of the tuple. See [LoadFromNpz]. - /// - /// E.g. for a two tuple (A, B) with `base == ""`, this will call: - /// 1. `self.0.read("0.", r)` - /// 2. `self.1.read("1.", r)` - fn read(&mut self, base: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - $(self.$idx.read(&format!("{}{}.", base, $idx), r)?;)+ - Ok(()) - } - } - /*This macro expands like this for a 4-tuple: impl< @@ -160,79 +130,6 @@ mod tests { assert!(model.1.bias.data() != m0.1.bias.data()); } - #[cfg(feature = "numpy")] - mod numpy_tests { - use super::*; - use std::fs::File; - use tempfile::NamedTempFile; - - #[test] - fn test_save_tuple() { - let model: ( - Linear<1, 2>, - ReLU, - Linear<2, 3>, - (Dropout, Linear<1, 2>, Linear<3, 4>), - ) = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - let mut names = zip.file_names().collect::>(); - names.sort_unstable(); - assert_eq!( - &names, - &[ - "0.bias.npy", - "0.weight.npy", - "2.bias.npy", - "2.weight.npy", - "3.1.bias.npy", - "3.1.weight.npy", - "3.2.bias.npy", - "3.2.weight.npy", - ] - ); - } - - #[test] - fn test_load_tuple() { - type Model = ( - Linear<1, 2>, - ReLU, - Linear<2, 3>, - (Dropout, Linear<1, 2>, Linear<3, 4>), - ); - - let mut rng = StdRng::seed_from_u64(0); - let mut saved_model: Model = Default::default(); - saved_model.reset_params(&mut rng); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded_model: Model = Default::default(); - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.0.weight.data(), saved_model.0.weight.data()); - assert_eq!(loaded_model.0.bias.data(), saved_model.0.bias.data()); - assert_eq!(loaded_model.2.weight.data(), saved_model.2.weight.data()); - assert_eq!(loaded_model.2.bias.data(), saved_model.2.bias.data()); - assert_eq!( - loaded_model.3 .1.weight.data(), - saved_model.3 .1.weight.data() - ); - assert_eq!(loaded_model.3 .1.bias.data(), saved_model.3 .1.bias.data()); - - assert_eq!( - loaded_model.3 .2.weight.data(), - saved_model.3 .2.weight.data() - ); - assert_eq!(loaded_model.3 .2.bias.data(), saved_model.3 .2.bias.data()); - } - } - /// A struct to test the forward method of tuples. This sets the `I`th valuein a 1d tensors of size `N` to 1.0. #[derive(Debug, Default, Clone)] struct SetTo1; diff --git a/src/nn/layer_norm.rs b/src/nn/layer_norm.rs index 6e98b1f4c..e52e0433d 100644 --- a/src/nn/layer_norm.rs +++ b/src/nn/layer_norm.rs @@ -2,10 +2,6 @@ use crate::arrays::Axis; use crate::devices::{Cpu, FillElements}; use crate::gradients::{CanUpdateWithGradients, GradientProvider, Tape, UnusedTensors}; use crate::prelude::*; -#[cfg(feature = "numpy")] -use std::io::{Read, Seek, Write}; -#[cfg(feature = "numpy")] -use zip::{result::ZipResult, ZipArchive}; /// Implements layer normalization as described in [Layer Normalization](https://arxiv.org/abs/1607.06450). /// @@ -116,28 +112,6 @@ where } } -#[cfg(feature = "numpy")] -impl SaveToNpz for LayerNorm1D { - /// Saves [Self::gamma] to `{pre}gamma.npy` and [Self::beta] to `{pre}beta.npy` - /// using [npz_fwrite()]. - fn write(&self, pre: &str, w: &mut zip::ZipWriter) -> ZipResult<()> { - npz_fwrite(w, format!("{pre}gamma.npy"), self.gamma.data())?; - npz_fwrite(w, format!("{pre}beta.npy"), self.beta.data())?; - Ok(()) - } -} - -#[cfg(feature = "numpy")] -impl LoadFromNpz for LayerNorm1D { - /// Reads [Self::gamma] from `{p}gamma.npy` and [Self::beta] from `{p}beta.npy` - /// using [npz_fread()]. - fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - npz_fread(r, format!("{p}gamma.npy"), self.gamma.mut_data())?; - npz_fread(r, format!("{p}beta.npy"), self.beta.mut_data())?; - Ok(()) - } -} - #[cfg(test)] mod tests { use super::*; @@ -204,62 +178,6 @@ mod tests { assert_close(g.ref_gradient(&m.beta), &[0.2; 5]); } - #[cfg(feature = "numpy")] - mod numpy_tests { - use super::*; - use std::fs::File; - use tempfile::NamedTempFile; - - #[test] - fn test_save_layer_norm() { - let model: LayerNorm1D<13> = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - let mut names = zip.file_names().collect::>(); - names.sort_unstable(); - assert_eq!(&names, &["beta.npy", "gamma.npy",]); - } - - #[test] - fn test_save_layer_norm_tuple() { - let model: (LayerNorm1D<5>, LayerNorm1D<13>) = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - let mut names = zip.file_names().collect::>(); - names.sort_unstable(); - assert_eq!( - &names, - &["0.beta.npy", "0.gamma.npy", "1.beta.npy", "1.gamma.npy"] - ); - } - - #[test] - fn test_load_layer_norm() { - let mut rng = StdRng::seed_from_u64(0); - let mut saved_model: LayerNorm1D<13> = Default::default(); - saved_model.gamma.randomize(&mut rng, &Standard); - saved_model.beta.randomize(&mut rng, &Standard); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded_model: LayerNorm1D<13> = Default::default(); - assert!(loaded_model.gamma.data() != saved_model.gamma.data()); - assert!(loaded_model.beta.data() != saved_model.beta.data()); - - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.gamma.data(), saved_model.gamma.data()); - assert_eq!(loaded_model.beta.data(), saved_model.beta.data()); - } - } #[test] fn test_layer_norm_missing_gradients() { let mut model: LayerNorm1D<5> = Default::default(); diff --git a/src/nn/linear.rs b/src/nn/linear.rs index c13ea544d..a9cf4fc01 100644 --- a/src/nn/linear.rs +++ b/src/nn/linear.rs @@ -2,10 +2,6 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, Tape, UnusedTen use crate::prelude::*; use rand::Rng; use rand_distr::Uniform; -#[cfg(feature = "numpy")] -use std::io::{Read, Seek, Write}; -#[cfg(feature = "numpy")] -use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// A linear transformation of the form `weight * x + bias`, where `weight` is a matrix, `x` is a vector or matrix, /// and `bias` is a vector. @@ -54,34 +50,6 @@ impl ResetParams for Linear { } } -#[cfg(feature = "numpy")] -impl SaveToNpz for Linear { - /// Saves [Self::weight] to `{pre}weight.npy` and [Self::bias] to `{pre}bias.npy` - /// using [npz_fwrite()]. - fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> - where - W: Write + Seek, - { - npz_fwrite(w, format!("{pre}weight.npy"), self.weight.data())?; - npz_fwrite(w, format!("{pre}bias.npy"), self.bias.data())?; - Ok(()) - } -} - -#[cfg(feature = "numpy")] -impl LoadFromNpz for Linear { - /// Reads [Self::weight] from `{pre}weight.npy` and [Self::bias] from `{pre}bias.npy` - /// using [npz_fread()]. - fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> - where - R: Read + Seek, - { - npz_fread(r, format!("{pre}weight.npy"), self.weight.mut_data())?; - npz_fread(r, format!("{pre}bias.npy"), self.bias.mut_data())?; - Ok(()) - } -} - impl Module> for Linear { type Output = Tensor1D; @@ -258,55 +226,6 @@ mod tests { ); } - #[cfg(feature = "numpy")] - mod numpy_tests { - use super::*; - use rand::{prelude::StdRng, SeedableRng}; - use std::fs::File; - use tempfile::NamedTempFile; - - #[test] - fn test_save_linear() { - let model: Linear<5, 3> = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let mut zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - { - let weight_file = zip - .by_name("weight.npy") - .expect("failed to find weight.npy file"); - assert!(weight_file.size() > 0); - } - { - let bias_file = zip - .by_name("bias.npy") - .expect("failed to find bias.npy file"); - assert!(bias_file.size() > 0); - } - } - - #[test] - fn test_load_linear() { - let mut rng = StdRng::seed_from_u64(0); - let mut saved_model: Linear<5, 3> = Default::default(); - saved_model.reset_params(&mut rng); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded_model: Linear<5, 3> = Default::default(); - assert!(loaded_model.weight.data() != saved_model.weight.data()); - assert!(loaded_model.bias.data() != saved_model.bias.data()); - - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.weight.data(), saved_model.weight.data()); - assert_eq!(loaded_model.bias.data(), saved_model.bias.data()); - } - } - #[test] fn test_linear_missing_gradients() { let mut model: Linear<5, 3> = Default::default(); diff --git a/src/nn/mod.rs b/src/nn/mod.rs index d4bc1b43b..ce60e39a1 100644 --- a/src/nn/mod.rs +++ b/src/nn/mod.rs @@ -108,6 +108,9 @@ mod npz; #[cfg(feature = "numpy")] pub use npz::*; +#[cfg(feature = "numpy")] +mod npz_impls; + #[cfg(feature = "nightly")] mod transformer; diff --git a/src/nn/npz_impls.rs b/src/nn/npz_impls.rs new file mode 100644 index 000000000..c6b934215 --- /dev/null +++ b/src/nn/npz_impls.rs @@ -0,0 +1,774 @@ +use super::npz::{npz_fread, npz_fwrite, LoadFromNpz, SaveToNpz}; +use crate::prelude::*; +use std::io::{Read, Seek, Write}; +use zip::{result::ZipResult, ZipArchive, ZipWriter}; + +impl SaveToNpz for BatchNorm2D { + fn write(&self, p: &str, w: &mut zip::ZipWriter) -> ZipResult<()> { + npz_fwrite(w, format!("{p}scale.npy"), self.scale.data())?; + npz_fwrite(w, format!("{p}bias.npy"), self.bias.data())?; + npz_fwrite(w, format!("{p}running_mean.npy"), self.running_mean.data())?; + npz_fwrite(w, format!("{p}running_var.npy"), self.running_var.data())?; + Ok(()) + } +} + +impl LoadFromNpz for BatchNorm2D { + fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + npz_fread(r, format!("{p}scale.npy"), self.scale.mut_data())?; + npz_fread(r, format!("{p}bias.npy"), self.bias.mut_data())?; + let mean = self.running_mean.mut_data(); + npz_fread(r, format!("{p}running_mean.npy"), mean)?; + let var = self.running_var.mut_data(); + npz_fread(r, format!("{p}running_var.npy"), var)?; + Ok(()) + } +} + +#[cfg(feature = "nightly")] +impl SaveToNpz + for Conv2D +{ + /// Saves [Self::weight] to `{pre}weight.npy` and [Self::bias] to `{pre}bias.npy` + /// using [npz_fwrite()]. + fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { + npz_fwrite(w, format!("{pre}weight.npy"), self.weight.data())?; + npz_fwrite(w, format!("{pre}bias.npy"), self.bias.data())?; + Ok(()) + } +} + +#[cfg(feature = "nightly")] +impl LoadFromNpz + for Conv2D +{ + /// Reads [Self::weight] from `{pre}weight.npy` and [Self::bias] from `{pre}bias.npy` + /// using [npz_fread()]. + fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + npz_fread(r, format!("{pre}weight.npy"), self.weight.mut_data())?; + npz_fread(r, format!("{pre}bias.npy"), self.bias.mut_data())?; + Ok(()) + } +} + +#[cfg(feature = "numpy")] +impl SaveToNpz for GeneralizedResidual { + /// Pass through to `F`/`R`'s [SaveToNpz]. + fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { + self.f.write(&format!("{p}.f"), w)?; + self.r.write(&format!("{p}.r"), w) + } +} + +#[cfg(feature = "numpy")] +impl LoadFromNpz for GeneralizedResidual { + /// Pass through to `F`/`R`'s [LoadFromNpz]. + fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + self.f.read(&format!("{p}.f"), r)?; + self.r.read(&format!("{p}.r"), r) + } +} + +#[cfg(feature = "numpy")] +impl SaveToNpz for LayerNorm1D { + /// Saves [Self::gamma] to `{pre}gamma.npy` and [Self::beta] to `{pre}beta.npy` + /// using [npz_fwrite()]. + fn write(&self, pre: &str, w: &mut zip::ZipWriter) -> ZipResult<()> { + npz_fwrite(w, format!("{pre}gamma.npy"), self.gamma.data())?; + npz_fwrite(w, format!("{pre}beta.npy"), self.beta.data())?; + Ok(()) + } +} + +#[cfg(feature = "numpy")] +impl LoadFromNpz for LayerNorm1D { + /// Reads [Self::gamma] from `{p}gamma.npy` and [Self::beta] from `{p}beta.npy` + /// using [npz_fread()]. + fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + npz_fread(r, format!("{p}gamma.npy"), self.gamma.mut_data())?; + npz_fread(r, format!("{p}beta.npy"), self.beta.mut_data())?; + Ok(()) + } +} + +impl SaveToNpz for Linear { + /// Saves [Self::weight] to `{pre}weight.npy` and [Self::bias] to `{pre}bias.npy` + /// using [npz_fwrite()]. + fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> + where + W: Write + Seek, + { + npz_fwrite(w, format!("{pre}weight.npy"), self.weight.data())?; + npz_fwrite(w, format!("{pre}bias.npy"), self.bias.data())?; + Ok(()) + } +} + +impl LoadFromNpz for Linear { + /// Reads [Self::weight] from `{pre}weight.npy` and [Self::bias] from `{pre}bias.npy` + /// using [npz_fread()]. + fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> + where + R: Read + Seek, + { + npz_fread(r, format!("{pre}weight.npy"), self.weight.mut_data())?; + npz_fread(r, format!("{pre}bias.npy"), self.bias.mut_data())?; + Ok(()) + } +} + +macro_rules! tuple_npz_impl { + ([$($name:ident),+], [$($idx:tt),+]) => { +impl<$($name: SaveToNpz),+> SaveToNpz for ($($name,)+) { + /// Calls `SaveToNpz::write(self., ...)` on each part of the tuple. See [SaveToNpz]. + /// + /// E.g. for a two tuple (A, B) with `base == ""`, this will call: + /// 1. `self.0.write("0.", w)` + /// 2. `self.1.write("1.", w)` + fn write(&self, base: &str, w: &mut ZipWriter) -> ZipResult<()> { + $(self.$idx.write(&format!("{}{}.", base, $idx), w)?;)+ + Ok(()) + } +} + +impl<$($name: LoadFromNpz),+> LoadFromNpz for ($($name,)+) { + /// Calls `LoadFromNpz::read(self., ...)` on each part of the tuple. See [LoadFromNpz]. + /// + /// E.g. for a two tuple (A, B) with `base == ""`, this will call: + /// 1. `self.0.read("0.", r)` + /// 2. `self.1.read("1.", r)` + fn read(&mut self, base: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + $(self.$idx.read(&format!("{}{}.", base, $idx), r)?;)+ + Ok(()) + } +} + }; +} + +tuple_npz_impl!([A, B], [0, 1]); +tuple_npz_impl!([A, B, C], [0, 1, 2]); +tuple_npz_impl!([A, B, C, D], [0, 1, 2, 3]); +tuple_npz_impl!([A, B, C, D, E], [0, 1, 2, 3, 4]); +tuple_npz_impl!([A, B, C, D, E, F], [0, 1, 2, 3, 4, 5]); + +#[cfg(feature = "numpy")] +impl SaveToNpz for Repeated { + /// Calls `SaveToNpz::write(self.modules[i], ...)` on each sub module. See [SaveToNpz]. + /// + /// E.g. for a two items with `base == ""`, this will call: + /// 1. `self.modules[0].write("0.", w)` + /// 2. `self.modules[1].write("1.", w)` + fn write(&self, base: &str, w: &mut ZipWriter) -> ZipResult<()> { + for i in 0..N { + self.modules[i].write(&format!("{base}{i}."), w)?; + } + Ok(()) + } +} + +#[cfg(feature = "numpy")] +impl LoadFromNpz for Repeated { + /// Calls `LoadFromNpz::read(self.modules[i], ...)` on each sub module. See [LoadFromNpz]. + /// + /// E.g. for a two items with `base == ""`, this will call: + /// 1. `self.modules[0].read("0.", r)` + /// 2. `self.modules[1].read("1.", r)` + fn read(&mut self, base: &str, r: &mut ZipArchive) -> Result<(), NpzError> + where + R: Read + Seek, + { + for i in 0..N { + self.modules[i].read(&format!("{base}{i}."), r)?; + } + Ok(()) + } +} + +#[cfg(feature = "numpy")] +impl SaveToNpz for Residual { + /// Pass through to `F`'s [SaveToNpz]. + fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { + self.0.write(p, w) + } +} + +#[cfg(feature = "numpy")] +impl LoadFromNpz for Residual { + /// Pass through to `F`'s [LoadFromNpz]. + fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + self.0.read(p, r) + } +} + +#[cfg(feature = "numpy")] +impl SaveToNpz for SplitInto { + fn write(&self, p: &str, w: &mut zip::ZipWriter) -> zip::result::ZipResult<()> + where + W: std::io::Write + std::io::Seek, + { + self.0.write(p, w) + } +} + +#[cfg(feature = "numpy")] +impl LoadFromNpz for SplitInto { + fn read(&mut self, p: &str, r: &mut zip::ZipArchive) -> Result<(), NpzError> + where + R: std::io::Read + std::io::Seek, + { + self.0.read(p, r) + } +} + +#[cfg(feature = "nightly")] +impl SaveToNpz + for TransformerDecoder +{ + fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { + self.0.write(pre, w) + } +} + +#[cfg(feature = "nightly")] +impl SaveToNpz + for TransformerDecoderBlock +{ + fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { + self.self_attn.write(&format!("{pre}self_attn."), w)?; + self.norm1.write(&format!("{pre}norm1."), w)?; + self.mh_attn.write(&format!("{pre}mh_attn."), w)?; + self.norm2.write(&format!("{pre}norm2."), w)?; + self.ff.0 .0.write(&format!("{pre}linear1."), w)?; + self.ff.0 .2.write(&format!("{pre}linear2."), w)?; + self.norm3.write(&format!("{pre}norm3."), w)?; + Ok(()) + } +} + +#[cfg(feature = "nightly")] +impl LoadFromNpz + for TransformerDecoderBlock +{ + fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + self.self_attn.read(&format!("{pre}self_attn."), r)?; + self.norm1.read(&format!("{pre}norm1."), r)?; + self.mh_attn.read(&format!("{pre}mh_attn."), r)?; + self.norm2.read(&format!("{pre}norm2."), r)?; + self.ff.0 .0.read(&format!("{pre}linear1."), r)?; + self.ff.0 .2.read(&format!("{pre}linear2."), r)?; + self.norm3.read(&format!("{pre}norm3."), r)?; + Ok(()) + } +} + +#[cfg(feature = "nightly")] +impl LoadFromNpz + for TransformerDecoder +{ + fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + self.0.read(pre, r) + } +} + +#[cfg(feature = "nightly")] +impl SaveToNpz + for TransformerEncoderBlock +{ + fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { + self.self_attn.write(&format!("{pre}self_attn."), w)?; + self.norm1.write(&format!("{pre}norm1."), w)?; + self.norm2.write(&format!("{pre}norm2."), w)?; + self.ff.0 .0.write(&format!("{pre}linear1."), w)?; + self.ff.0 .2.write(&format!("{pre}linear2."), w)?; + Ok(()) + } +} + +#[cfg(feature = "nightly")] +impl LoadFromNpz + for TransformerEncoderBlock +{ + fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + self.self_attn.read(&format!("{pre}self_attn."), r)?; + self.norm1.read(&format!("{pre}norm1."), r)?; + self.norm2.read(&format!("{pre}norm2."), r)?; + self.ff.0 .0.read(&format!("{pre}linear1."), r)?; + self.ff.0 .2.read(&format!("{pre}linear2."), r)?; + Ok(()) + } +} + +#[cfg(feature = "nightly")] +impl SaveToNpz + for MultiHeadAttention +{ + fn write(&self, pre: &str, w: &mut zip::ZipWriter) -> zip::result::ZipResult<()> + where + W: std::io::Write + std::io::Seek, + { + self.w_q.write(&format!("{pre}w_q."), w)?; + self.w_k.write(&format!("{pre}w_k."), w)?; + self.w_v.write(&format!("{pre}w_v."), w)?; + self.w_o.write(&format!("{pre}w_o."), w)?; + Ok(()) + } +} + +#[cfg(feature = "nightly")] +impl LoadFromNpz + for MultiHeadAttention +{ + fn read(&mut self, pre: &str, r: &mut zip::ZipArchive) -> Result<(), NpzError> + where + R: std::io::Read + std::io::Seek, + { + self.w_q.read(&format!("{pre}w_q."), r)?; + self.w_k.read(&format!("{pre}w_k."), r)?; + self.w_v.read(&format!("{pre}w_v."), r)?; + self.w_o.read(&format!("{pre}w_o."), r)?; + Ok(()) + } +} + +#[cfg(feature = "nightly")] +impl SaveToNpz + for Transformer +{ + fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { + self.encoder.write(&format!("{pre}encoder."), w)?; + self.decoder.write(&format!("{pre}decoder."), w)?; + Ok(()) + } +} + +#[cfg(feature = "nightly")] +impl LoadFromNpz + for Transformer +{ + fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + self.encoder.read(&format!("{pre}encoder."), r)?; + self.decoder.read(&format!("{pre}decoder."), r)?; + Ok(()) + } +} + +macro_rules! empty_npz_impl { + ($TyName:ty) => { + impl SaveToNpz for $TyName {} + impl LoadFromNpz for $TyName {} + }; +} + +empty_npz_impl!(ReLU); +empty_npz_impl!(Sin); +empty_npz_impl!(Cos); +empty_npz_impl!(Ln); +empty_npz_impl!(Exp); +empty_npz_impl!(Sigmoid); +empty_npz_impl!(Tanh); +empty_npz_impl!(Square); +empty_npz_impl!(Sqrt); +empty_npz_impl!(Abs); +empty_npz_impl!(Softmax); +empty_npz_impl!(Dropout); +empty_npz_impl!(AvgPoolGlobal); +empty_npz_impl!(MaxPoolGlobal); +empty_npz_impl!(MinPoolGlobal); +#[cfg(feature = "nightly")] +empty_npz_impl!(FlattenImage); + +impl SaveToNpz for DropoutOneIn {} +impl LoadFromNpz for DropoutOneIn {} + +#[cfg(feature = "nightly")] +impl SaveToNpz for AvgPool2D {} +#[cfg(feature = "nightly")] +impl LoadFromNpz for AvgPool2D {} + +#[cfg(feature = "nightly")] +impl SaveToNpz for MaxPool2D {} +#[cfg(feature = "nightly")] +impl LoadFromNpz for MaxPool2D {} + +#[cfg(feature = "nightly")] +impl SaveToNpz for MinPool2D {} +#[cfg(feature = "nightly")] +impl LoadFromNpz for MinPool2D {} + +#[cfg(test)] +mod tests { + use super::*; + use crate::gradients::*; + use crate::nn::tests::SimpleGradients; + use rand::thread_rng; + use rand_distr::Standard; + use std::fs::File; + use tempfile::NamedTempFile; + + #[test] + fn test_batchnorm2d_save_load() { + let mut rng = thread_rng(); + let mut bn: BatchNorm2D<3> = Default::default(); + + assert_eq!(bn.running_mean.data(), &[0.0; 3]); + assert_eq!(bn.running_var.data(), &[1.0; 3]); + assert_eq!(bn.scale.data(), &[1.0; 3]); + assert_eq!(bn.bias.data(), &[0.0; 3]); + + let x1: Tensor3D<3, 4, 5> = TensorCreator::randn(&mut rng); + let g = backward(bn.forward_mut(x1.trace()).exp().mean()); + bn.update(&mut SimpleGradients(g), &mut Default::default()); + + assert_ne!(bn.running_mean.data(), &[0.0; 3]); + assert_ne!(bn.running_var.data(), &[1.0; 3]); + assert_ne!(bn.scale.data(), &[1.0; 3]); + assert_ne!(bn.bias.data(), &[0.0; 3]); + + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(bn.save(file.path().to_str().unwrap()).is_ok()); + + let mut loaded: BatchNorm2D<3> = Default::default(); + assert!(loaded.load(file.path().to_str().unwrap()).is_ok()); + assert_eq!(loaded.scale.data(), bn.scale.data()); + assert_eq!(loaded.bias.data(), bn.bias.data()); + assert_eq!(loaded.running_mean.data(), bn.running_mean.data()); + assert_eq!(loaded.running_var.data(), bn.running_var.data()); + } + + #[cfg(feature = "nightly")] + #[test] + fn test_save_conv2d() { + let model: Conv2D<2, 4, 3> = Default::default(); + let file = NamedTempFile::new().expect("failed to create tempfile"); + model + .save(file.path().to_str().unwrap()) + .expect("failed to save model"); + let f = File::open(file.path()).expect("failed to open resulting file"); + let mut zip = ZipArchive::new(f).expect("failed to create zip archive from file"); + { + let weight_file = zip + .by_name("weight.npy") + .expect("failed to find weight.npy file"); + assert!(weight_file.size() > 0); + } + { + let bias_file = zip + .by_name("bias.npy") + .expect("failed to find bias.npy file"); + assert!(bias_file.size() > 0); + } + } + + #[cfg(feature = "nightly")] + #[test] + fn test_load_conv() { + let mut rng = thread_rng(); + let mut saved_model: Conv2D<2, 4, 3> = Default::default(); + saved_model.reset_params(&mut rng); + + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + + let mut loaded_model: Conv2D<2, 4, 3> = Default::default(); + assert!(loaded_model.weight.data() != saved_model.weight.data()); + assert!(loaded_model.bias.data() != saved_model.bias.data()); + + assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); + assert_eq!(loaded_model.weight.data(), saved_model.weight.data()); + assert_eq!(loaded_model.bias.data(), saved_model.bias.data()); + } + + #[test] + fn test_save_load_generalized_residual() { + let mut rng = thread_rng(); + let mut saved_model: GeneralizedResidual, Linear<5, 3>> = Default::default(); + saved_model.reset_params(&mut rng); + + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + + let mut loaded_model: GeneralizedResidual, Linear<5, 3>> = Default::default(); + assert_ne!(loaded_model.f.weight.data(), saved_model.f.weight.data()); + assert_ne!(loaded_model.f.bias.data(), saved_model.f.bias.data()); + assert_ne!(loaded_model.r.weight.data(), saved_model.r.weight.data()); + assert_ne!(loaded_model.r.bias.data(), saved_model.r.bias.data()); + + assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); + assert_eq!(loaded_model.f.weight.data(), saved_model.f.weight.data()); + assert_eq!(loaded_model.f.bias.data(), saved_model.f.bias.data()); + assert_eq!(loaded_model.r.weight.data(), saved_model.r.weight.data()); + assert_eq!(loaded_model.r.bias.data(), saved_model.r.bias.data()); + } + + #[test] + fn test_save_linear() { + let model: Linear<5, 3> = Default::default(); + let file = NamedTempFile::new().expect("failed to create tempfile"); + model + .save(file.path().to_str().unwrap()) + .expect("failed to save model"); + let f = File::open(file.path()).expect("failed to open resulting file"); + let mut zip = ZipArchive::new(f).expect("failed to create zip archive from file"); + { + let weight_file = zip + .by_name("weight.npy") + .expect("failed to find weight.npy file"); + assert!(weight_file.size() > 0); + } + { + let bias_file = zip + .by_name("bias.npy") + .expect("failed to find bias.npy file"); + assert!(bias_file.size() > 0); + } + } + + #[test] + fn test_load_linear() { + let mut rng = thread_rng(); + let mut saved_model: Linear<5, 3> = Default::default(); + saved_model.reset_params(&mut rng); + + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + + let mut loaded_model: Linear<5, 3> = Default::default(); + assert!(loaded_model.weight.data() != saved_model.weight.data()); + assert!(loaded_model.bias.data() != saved_model.bias.data()); + + assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); + assert_eq!(loaded_model.weight.data(), saved_model.weight.data()); + assert_eq!(loaded_model.bias.data(), saved_model.bias.data()); + } + + #[test] + fn test_save_tuple() { + let model: ( + Linear<1, 2>, + ReLU, + Linear<2, 3>, + (Dropout, Linear<1, 2>, Linear<3, 4>), + ) = Default::default(); + let file = NamedTempFile::new().expect("failed to create tempfile"); + model + .save(file.path().to_str().unwrap()) + .expect("failed to save model"); + let f = File::open(file.path()).expect("failed to open resulting file"); + let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); + let mut names = zip.file_names().collect::>(); + names.sort_unstable(); + assert_eq!( + &names, + &[ + "0.bias.npy", + "0.weight.npy", + "2.bias.npy", + "2.weight.npy", + "3.1.bias.npy", + "3.1.weight.npy", + "3.2.bias.npy", + "3.2.weight.npy", + ] + ); + } + + #[test] + fn test_load_tuple() { + type Model = ( + Linear<1, 2>, + ReLU, + Linear<2, 3>, + (Dropout, Linear<1, 2>, Linear<3, 4>), + ); + + let mut rng = thread_rng(); + let mut saved_model: Model = Default::default(); + saved_model.reset_params(&mut rng); + + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + + let mut loaded_model: Model = Default::default(); + assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); + assert_eq!(loaded_model.0.weight.data(), saved_model.0.weight.data()); + assert_eq!(loaded_model.0.bias.data(), saved_model.0.bias.data()); + assert_eq!(loaded_model.2.weight.data(), saved_model.2.weight.data()); + assert_eq!(loaded_model.2.bias.data(), saved_model.2.bias.data()); + assert_eq!( + loaded_model.3 .1.weight.data(), + saved_model.3 .1.weight.data() + ); + assert_eq!(loaded_model.3 .1.bias.data(), saved_model.3 .1.bias.data()); + + assert_eq!( + loaded_model.3 .2.weight.data(), + saved_model.3 .2.weight.data() + ); + assert_eq!(loaded_model.3 .2.bias.data(), saved_model.3 .2.bias.data()); + } + + #[test] + fn test_save_layer_norm() { + let model: LayerNorm1D<13> = Default::default(); + let file = NamedTempFile::new().expect("failed to create tempfile"); + model + .save(file.path().to_str().unwrap()) + .expect("failed to save model"); + let f = File::open(file.path()).expect("failed to open resulting file"); + let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); + let mut names = zip.file_names().collect::>(); + names.sort_unstable(); + assert_eq!(&names, &["beta.npy", "gamma.npy",]); + } + + #[test] + fn test_save_layer_norm_tuple() { + let model: (LayerNorm1D<5>, LayerNorm1D<13>) = Default::default(); + let file = NamedTempFile::new().expect("failed to create tempfile"); + model + .save(file.path().to_str().unwrap()) + .expect("failed to save model"); + let f = File::open(file.path()).expect("failed to open resulting file"); + let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); + let mut names = zip.file_names().collect::>(); + names.sort_unstable(); + assert_eq!( + &names, + &["0.beta.npy", "0.gamma.npy", "1.beta.npy", "1.gamma.npy"] + ); + } + + #[test] + fn test_load_layer_norm() { + let mut rng = thread_rng(); + let mut saved_model: LayerNorm1D<13> = Default::default(); + saved_model.gamma.randomize(&mut rng, &Standard); + saved_model.beta.randomize(&mut rng, &Standard); + + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + + let mut loaded_model: LayerNorm1D<13> = Default::default(); + assert!(loaded_model.gamma.data() != saved_model.gamma.data()); + assert!(loaded_model.beta.data() != saved_model.beta.data()); + + assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); + assert_eq!(loaded_model.gamma.data(), saved_model.gamma.data()); + assert_eq!(loaded_model.beta.data(), saved_model.beta.data()); + } + + #[test] + fn test_save_repeated() { + let model: Repeated, 4> = Default::default(); + let file = NamedTempFile::new().expect("failed to create tempfile"); + model + .save(file.path().to_str().unwrap()) + .expect("failed to save model"); + let f = File::open(file.path()).expect("failed to open resulting file"); + let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); + let mut names = zip.file_names().collect::>(); + names.sort_unstable(); + assert_eq!( + &names, + &[ + "0.bias.npy", + "0.weight.npy", + "1.bias.npy", + "1.weight.npy", + "2.bias.npy", + "2.weight.npy", + "3.bias.npy", + "3.weight.npy", + ] + ); + } + + #[test] + fn test_load_repeated() { + type Model = Repeated, 4>; + + let mut rng = thread_rng(); + let mut saved_model: Model = Default::default(); + saved_model.reset_params(&mut rng); + + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + + let mut loaded_model: Model = Default::default(); + assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); + for i in 0..4 { + assert_eq!( + loaded_model.modules[i].weight.data(), + saved_model.modules[i].weight.data() + ); + assert_eq!( + loaded_model.modules[i].bias.data(), + saved_model.modules[i].bias.data() + ); + } + } + + #[test] + fn test_save_load_residual() { + let mut rng = thread_rng(); + let mut saved_model: Residual> = Default::default(); + saved_model.reset_params(&mut rng); + + let file = NamedTempFile::new().expect("failed to create tempfile"); + assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + + let mut loaded_model: Residual> = Default::default(); + assert_ne!(loaded_model.0.weight.data(), saved_model.0.weight.data()); + assert_ne!(loaded_model.0.bias.data(), saved_model.0.bias.data()); + + assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); + assert_eq!(loaded_model.0.weight.data(), saved_model.0.weight.data()); + assert_eq!(loaded_model.0.bias.data(), saved_model.0.bias.data()); + } + + #[cfg(feature = "nightly")] + #[test] + fn test_save_and_load() { + let mut rng = thread_rng(); + + let mut saved: MultiHeadAttention<12, 4> = Default::default(); + saved.reset_params(&mut rng); + + let file = NamedTempFile::new().expect("failed to create tempfile"); + saved.save(file.path()).expect(""); + + let mut loaded: MultiHeadAttention<12, 4> = Default::default(); + loaded.load(file.path()).expect(""); + + let q: Tensor3D<2, 3, 12> = TensorCreator::randn(&mut rng); + let k: Tensor3D<2, 4, 12> = TensorCreator::randn(&mut rng); + let v: Tensor3D<2, 4, 12> = TensorCreator::randn(&mut rng); + let y1: Tensor3D<2, 3, 12, _> = saved.forward((q.clone(), k.clone(), v.clone())); + let y2: Tensor3D<2, 3, 12, _> = loaded.forward((q.clone(), k.clone(), v.clone())); + + assert_eq!(y1.data(), y2.data()); + } + + #[cfg(feature = "nightly")] + #[test] + fn test_save_load() { + let mut rng = thread_rng(); + + let mut saved: Transformer<16, 4, 3, 4, 8> = Default::default(); + saved.reset_params(&mut rng); + + let file = NamedTempFile::new().expect("failed to create tempfile"); + saved.save(file.path()).expect(""); + + let mut loaded: Transformer<16, 4, 3, 4, 8> = Default::default(); + loaded.load(file.path()).expect(""); + + let src: Tensor3D<4, 12, 16> = TensorCreator::randn(&mut rng); + let tgt: Tensor3D<4, 6, 16> = TensorCreator::randn(&mut rng); + + let y1 = saved.forward_mut((src.clone(), tgt.clone())); + let y2 = loaded.forward_mut((src.clone(), tgt.clone())); + + assert_eq!(y1.data(), y2.data()); + } +} diff --git a/src/nn/pool2d.rs b/src/nn/pool2d.rs index 9fda878d6..3acbd1fab 100644 --- a/src/nn/pool2d.rs +++ b/src/nn/pool2d.rs @@ -1,5 +1,3 @@ -#[cfg(feature = "numpy")] -use super::{LoadFromNpz, SaveToNpz}; use super::{Module, ModuleMut, ResetParams}; use crate::gradients::*; use crate::tensor::*; @@ -47,11 +45,6 @@ macro_rules! impl_pools { fn reset_params(&mut self, _: &mut R) {} } - #[cfg(feature = "numpy")] - impl SaveToNpz for $PoolTy {} - #[cfg(feature = "numpy")] - impl LoadFromNpz for $PoolTy {} - impl< const K: usize, const S: usize, diff --git a/src/nn/pool_global.rs b/src/nn/pool_global.rs index 233139966..0db2c365a 100644 --- a/src/nn/pool_global.rs +++ b/src/nn/pool_global.rs @@ -1,5 +1,3 @@ -#[cfg(feature = "numpy")] -use super::{LoadFromNpz, SaveToNpz}; use super::{Module, ModuleMut, ResetParams}; use crate::gradients::*; use crate::tensor::*; @@ -67,12 +65,6 @@ macro_rules! impl_pools { fn update(&mut self, _: &mut G, _: &mut UnusedTensors) {} } - #[cfg(feature = "numpy")] - impl SaveToNpz for $PoolTy {} - - #[cfg(feature = "numpy")] - impl LoadFromNpz for $PoolTy {} - impl Module> for $PoolTy { type Output = Tensor1D; fn forward(&self, input: Tensor2D) -> Self::Output { diff --git a/src/nn/repeated.rs b/src/nn/repeated.rs index d75894718..43e42f768 100644 --- a/src/nn/repeated.rs +++ b/src/nn/repeated.rs @@ -1,9 +1,5 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, UnusedTensors}; use crate::prelude::*; -#[cfg(feature = "numpy")] -use std::io::{Read, Seek, Write}; -#[cfg(feature = "numpy")] -use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// Repeats `T` `N` times. This requires that `T`'s input is the same as it's output. /// @@ -56,39 +52,6 @@ impl CanUpdateWithGradients for Repea } } -#[cfg(feature = "numpy")] -impl SaveToNpz for Repeated { - /// Calls `SaveToNpz::write(self.modules[i], ...)` on each sub module. See [SaveToNpz]. - /// - /// E.g. for a two items with `base == ""`, this will call: - /// 1. `self.modules[0].write("0.", w)` - /// 2. `self.modules[1].write("1.", w)` - fn write(&self, base: &str, w: &mut ZipWriter) -> ZipResult<()> { - for i in 0..N { - self.modules[i].write(&format!("{base}{i}."), w)?; - } - Ok(()) - } -} - -#[cfg(feature = "numpy")] -impl LoadFromNpz for Repeated { - /// Calls `LoadFromNpz::read(self.modules[i], ...)` on each sub module. See [LoadFromNpz]. - /// - /// E.g. for a two items with `base == ""`, this will call: - /// 1. `self.modules[0].read("0.", r)` - /// 2. `self.modules[1].read("1.", r)` - fn read(&mut self, base: &str, r: &mut ZipArchive) -> Result<(), NpzError> - where - R: Read + Seek, - { - for i in 0..N { - self.modules[i].read(&format!("{base}{i}."), r)?; - } - Ok(()) - } -} - impl, const N: usize> Module for Repeated { type Output = T::Output; fn forward(&self, mut x: Input) -> Self::Output { @@ -156,64 +119,6 @@ mod tests { assert_eq!(x.data(), m.forward_mut(Tensor1D::zeros()).data()); } - #[cfg(feature = "numpy")] - mod numpy_tests { - use super::*; - use std::fs::File; - use tempfile::NamedTempFile; - - #[test] - fn test_save_repeated() { - let model: Repeated, 4> = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - let mut names = zip.file_names().collect::>(); - names.sort_unstable(); - assert_eq!( - &names, - &[ - "0.bias.npy", - "0.weight.npy", - "1.bias.npy", - "1.weight.npy", - "2.bias.npy", - "2.weight.npy", - "3.bias.npy", - "3.weight.npy", - ] - ); - } - - #[test] - fn test_load_repeated() { - type Model = Repeated, 4>; - - let mut rng = StdRng::seed_from_u64(0); - let mut saved_model: Model = Default::default(); - saved_model.reset_params(&mut rng); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded_model: Model = Default::default(); - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - for i in 0..4 { - assert_eq!( - loaded_model.modules[i].weight.data(), - saved_model.modules[i].weight.data() - ); - assert_eq!( - loaded_model.modules[i].bias.data(), - saved_model.modules[i].bias.data() - ); - } - } - } - #[test] fn test_repeated_missing_gradients() { let mut model: Repeated, 3> = Default::default(); diff --git a/src/nn/residual.rs b/src/nn/residual.rs index cb3840576..939ed088f 100644 --- a/src/nn/residual.rs +++ b/src/nn/residual.rs @@ -1,9 +1,5 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, UnusedTensors}; use crate::prelude::*; -#[cfg(feature = "numpy")] -use std::io::{Read, Seek, Write}; -#[cfg(feature = "numpy")] -use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// A residual connection around `F`: `F(x) + x`, /// as introduced in [Deep Residual Learning for Image Recognition](https://arxiv.org/abs/1512.03385). @@ -52,22 +48,6 @@ impl, F: ModuleMut> ModuleMut for Resid } } -#[cfg(feature = "numpy")] -impl SaveToNpz for Residual { - /// Pass through to `F`'s [SaveToNpz]. - fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { - self.0.write(p, w) - } -} - -#[cfg(feature = "numpy")] -impl LoadFromNpz for Residual { - /// Pass through to `F`'s [LoadFromNpz]. - fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - self.0.read(p, r) - } -} - #[cfg(test)] mod tests { use super::*; @@ -104,28 +84,4 @@ mod tests { assert_close(g.ref_gradient(&model.0.bias), &[0.5; 2]); assert_close(g.ref_gradient(&x), &[[0.18806472, 0.21419683]; 4]); } - - #[cfg(feature = "numpy")] - mod numpy_tests { - use super::*; - use tempfile::NamedTempFile; - - #[test] - fn test_save_load_residual() { - let mut rng = StdRng::seed_from_u64(0); - let mut saved_model: Residual> = Default::default(); - saved_model.reset_params(&mut rng); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded_model: Residual> = Default::default(); - assert_ne!(loaded_model.0.weight.data(), saved_model.0.weight.data()); - assert_ne!(loaded_model.0.bias.data(), saved_model.0.bias.data()); - - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.0.weight.data(), saved_model.0.weight.data()); - assert_eq!(loaded_model.0.bias.data(), saved_model.0.bias.data()); - } - } } diff --git a/src/nn/split_into.rs b/src/nn/split_into.rs index a542ee210..1db0af88e 100644 --- a/src/nn/split_into.rs +++ b/src/nn/split_into.rs @@ -32,26 +32,6 @@ impl ResetParams for SplitInto { } } -#[cfg(feature = "numpy")] -impl SaveToNpz for SplitInto { - fn write(&self, p: &str, w: &mut zip::ZipWriter) -> zip::result::ZipResult<()> - where - W: std::io::Write + std::io::Seek, - { - self.0.write(p, w) - } -} - -#[cfg(feature = "numpy")] -impl LoadFromNpz for SplitInto { - fn read(&mut self, p: &str, r: &mut zip::ZipArchive) -> Result<(), NpzError> - where - R: std::io::Read + std::io::Seek, - { - self.0.read(p, r) - } -} - macro_rules! tuple_impls { ([$($heads:ident),+] $tail:ident) => { impl< diff --git a/src/nn/transformer/decoder.rs b/src/nn/transformer/decoder.rs index 0373712c0..96a4a5961 100644 --- a/src/nn/transformer/decoder.rs +++ b/src/nn/transformer/decoder.rs @@ -1,10 +1,6 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, UnusedTensors}; use crate::prelude::*; use rand::Rng; -#[cfg(feature = "numpy")] -use std::io::{Read, Seek, Write}; -#[cfg(feature = "numpy")] -use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// **Requires Nightly** A transformer decoder. /// @@ -67,24 +63,6 @@ where } } -#[cfg(feature = "numpy")] -impl SaveToNpz - for TransformerDecoder -{ - fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { - self.0.write(pre, w) - } -} - -#[cfg(feature = "numpy")] -impl LoadFromNpz - for TransformerDecoder -{ - fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - self.0.read(pre, r) - } -} - /// **Requires Nightly** A transformer decoder block. Different than the normal transformer block /// as this self attention accepts an additional sequence from the encoder. /// @@ -173,38 +151,6 @@ where } } -#[cfg(feature = "numpy")] -impl SaveToNpz - for TransformerDecoderBlock -{ - fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { - self.self_attn.write(&format!("{pre}self_attn."), w)?; - self.norm1.write(&format!("{pre}norm1."), w)?; - self.mh_attn.write(&format!("{pre}mh_attn."), w)?; - self.norm2.write(&format!("{pre}norm2."), w)?; - self.ff.0 .0.write(&format!("{pre}linear1."), w)?; - self.ff.0 .2.write(&format!("{pre}linear2."), w)?; - self.norm3.write(&format!("{pre}norm3."), w)?; - Ok(()) - } -} - -#[cfg(feature = "numpy")] -impl LoadFromNpz - for TransformerDecoderBlock -{ - fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - self.self_attn.read(&format!("{pre}self_attn."), r)?; - self.norm1.read(&format!("{pre}norm1."), r)?; - self.mh_attn.read(&format!("{pre}mh_attn."), r)?; - self.norm2.read(&format!("{pre}norm2."), r)?; - self.ff.0 .0.read(&format!("{pre}linear1."), r)?; - self.ff.0 .2.read(&format!("{pre}linear2."), r)?; - self.norm3.read(&format!("{pre}norm3."), r)?; - Ok(()) - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/nn/transformer/encoder.rs b/src/nn/transformer/encoder.rs index c8273509f..7ae60cdb1 100644 --- a/src/nn/transformer/encoder.rs +++ b/src/nn/transformer/encoder.rs @@ -1,9 +1,5 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, UnusedTensors}; use crate::prelude::*; -#[cfg(feature = "numpy")] -use std::io::{Read, Seek, Write}; -#[cfg(feature = "numpy")] -use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// **Requires Nightly** A transformer encoder. /// @@ -41,10 +37,10 @@ pub struct TransformerEncoderBlock< const NUM_HEADS: usize, const FF_DIM: usize, > { - self_attn: MultiHeadAttention, - norm1: LayerNorm1D, - ff: FF, - norm2: LayerNorm1D, + pub self_attn: MultiHeadAttention, + pub norm1: LayerNorm1D, + pub ff: FF, + pub norm2: LayerNorm1D, } type FF = Residual<(Linear, ReLU, Linear)>; @@ -107,34 +103,6 @@ where } } -#[cfg(feature = "numpy")] -impl SaveToNpz - for TransformerEncoderBlock -{ - fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { - self.self_attn.write(&format!("{pre}self_attn."), w)?; - self.norm1.write(&format!("{pre}norm1."), w)?; - self.norm2.write(&format!("{pre}norm2."), w)?; - self.ff.0 .0.write(&format!("{pre}linear1."), w)?; - self.ff.0 .2.write(&format!("{pre}linear2."), w)?; - Ok(()) - } -} - -#[cfg(feature = "numpy")] -impl LoadFromNpz - for TransformerEncoderBlock -{ - fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - self.self_attn.read(&format!("{pre}self_attn."), r)?; - self.norm1.read(&format!("{pre}norm1."), r)?; - self.norm2.read(&format!("{pre}norm2."), r)?; - self.ff.0 .0.read(&format!("{pre}linear1."), r)?; - self.ff.0 .2.read(&format!("{pre}linear2."), r)?; - Ok(()) - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/nn/transformer/mha.rs b/src/nn/transformer/mha.rs index 1c55a4538..21ffdad28 100644 --- a/src/nn/transformer/mha.rs +++ b/src/nn/transformer/mha.rs @@ -53,38 +53,6 @@ impl CanUpdateWi } } -#[cfg(feature = "numpy")] -impl SaveToNpz - for MultiHeadAttention -{ - fn write(&self, pre: &str, w: &mut zip::ZipWriter) -> zip::result::ZipResult<()> - where - W: std::io::Write + std::io::Seek, - { - self.w_q.write(&format!("{pre}w_q."), w)?; - self.w_k.write(&format!("{pre}w_k."), w)?; - self.w_v.write(&format!("{pre}w_v."), w)?; - self.w_o.write(&format!("{pre}w_o."), w)?; - Ok(()) - } -} - -#[cfg(feature = "numpy")] -impl LoadFromNpz - for MultiHeadAttention -{ - fn read(&mut self, pre: &str, r: &mut zip::ZipArchive) -> Result<(), NpzError> - where - R: std::io::Read + std::io::Seek, - { - self.w_q.read(&format!("{pre}w_q."), r)?; - self.w_k.read(&format!("{pre}w_k."), r)?; - self.w_v.read(&format!("{pre}w_v."), r)?; - self.w_o.read(&format!("{pre}w_o."), r)?; - Ok(()) - } -} - impl< const M: usize, const H: usize, @@ -217,7 +185,6 @@ mod tests { use super::*; use crate::{nn::tests::SimpleGradients, tests::assert_close}; use rand::{rngs::StdRng, thread_rng, SeedableRng}; - use tempfile::NamedTempFile; #[test] fn test_mha_unbatched() { @@ -327,27 +294,4 @@ mod tests { mha.update(&mut g, &mut unused); assert!(unused.is_empty()); } - - #[cfg(feature = "numpy")] - #[test] - fn test_save_and_load() { - let mut rng = thread_rng(); - - let mut saved: MultiHeadAttention<12, 4> = Default::default(); - saved.reset_params(&mut rng); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - saved.save(file.path()).expect(""); - - let mut loaded: MultiHeadAttention<12, 4> = Default::default(); - loaded.load(file.path()).expect(""); - - let q: Tensor3D<2, 3, 12> = TensorCreator::randn(&mut rng); - let k: Tensor3D<2, 4, 12> = TensorCreator::randn(&mut rng); - let v: Tensor3D<2, 4, 12> = TensorCreator::randn(&mut rng); - let y1: Tensor3D<2, 3, 12, _> = saved.forward((q.clone(), k.clone(), v.clone())); - let y2: Tensor3D<2, 3, 12, _> = loaded.forward((q.clone(), k.clone(), v.clone())); - - assert_eq!(y1.data(), y2.data()); - } } diff --git a/src/nn/transformer/transformer.rs b/src/nn/transformer/transformer.rs index 74206d46d..c56883ce6 100644 --- a/src/nn/transformer/transformer.rs +++ b/src/nn/transformer/transformer.rs @@ -1,9 +1,5 @@ use crate::gradients::{CanUpdateWithGradients, GradientProvider, UnusedTensors}; use crate::prelude::*; -#[cfg(feature = "numpy")] -use std::io::{Read, Seek, Write}; -#[cfg(feature = "numpy")] -use zip::{result::ZipResult, ZipArchive, ZipWriter}; /// **Requires Nightly** Transformer architecture as described in /// [Attention is all you need](https://arxiv.org/abs/1706.03762). @@ -36,8 +32,8 @@ pub struct Transformer< const NUM_DECODER_LAYERS: usize, const FF_DIM: usize, > { - encoder: TransformerEncoder, - decoder: TransformerDecoder, + pub encoder: TransformerEncoder, + pub decoder: TransformerDecoder, } impl ResetParams @@ -88,34 +84,11 @@ where } } -#[cfg(feature = "numpy")] -impl SaveToNpz - for Transformer -{ - fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { - self.encoder.write(&format!("{pre}encoder."), w)?; - self.decoder.write(&format!("{pre}decoder."), w)?; - Ok(()) - } -} - -#[cfg(feature = "numpy")] -impl LoadFromNpz - for Transformer -{ - fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - self.encoder.read(&format!("{pre}encoder."), r)?; - self.decoder.read(&format!("{pre}decoder."), r)?; - Ok(()) - } -} - #[cfg(test)] mod tests { use super::*; use crate::nn::tests::SimpleGradients; - use rand::{rngs::StdRng, thread_rng, SeedableRng}; - use tempfile::NamedTempFile; + use rand::{rngs::StdRng, SeedableRng}; #[test] fn test_forward() { @@ -151,27 +124,4 @@ mod tests { assert!(unused.is_empty()); } - - #[cfg(feature = "numpy")] - #[test] - fn test_save_load() { - let mut rng = thread_rng(); - - let mut saved: Transformer<16, 4, 3, 4, 8> = Default::default(); - saved.reset_params(&mut rng); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - saved.save(file.path()).expect(""); - - let mut loaded: Transformer<16, 4, 3, 4, 8> = Default::default(); - loaded.load(file.path()).expect(""); - - let src: Tensor3D<4, 12, 16> = TensorCreator::randn(&mut rng); - let tgt: Tensor3D<4, 6, 16> = TensorCreator::randn(&mut rng); - - let y1 = saved.forward_mut((src.clone(), tgt.clone())); - let y2 = loaded.forward_mut((src.clone(), tgt.clone())); - - assert_eq!(y1.data(), y2.data()); - } } From 84415d2a387e68e857526db66c65364e49852214 Mon Sep 17 00:00:00 2001 From: Corey Lowman Date: Sat, 15 Oct 2022 12:11:38 -0400 Subject: [PATCH 4/6] Clean up npz impls --- src/nn/npz_impls.rs | 185 ++++++++++++++------------------------------ 1 file changed, 59 insertions(+), 126 deletions(-) diff --git a/src/nn/npz_impls.rs b/src/nn/npz_impls.rs index c6b934215..41995dd2d 100644 --- a/src/nn/npz_impls.rs +++ b/src/nn/npz_impls.rs @@ -29,11 +29,9 @@ impl LoadFromNpz for BatchNorm2D { impl SaveToNpz for Conv2D { - /// Saves [Self::weight] to `{pre}weight.npy` and [Self::bias] to `{pre}bias.npy` - /// using [npz_fwrite()]. - fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { - npz_fwrite(w, format!("{pre}weight.npy"), self.weight.data())?; - npz_fwrite(w, format!("{pre}bias.npy"), self.bias.data())?; + fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { + npz_fwrite(w, format!("{p}weight.npy"), self.weight.data())?; + npz_fwrite(w, format!("{p}bias.npy"), self.bias.data())?; Ok(()) } } @@ -42,48 +40,36 @@ impl LoadFromNpz for Conv2D { - /// Reads [Self::weight] from `{pre}weight.npy` and [Self::bias] from `{pre}bias.npy` - /// using [npz_fread()]. - fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - npz_fread(r, format!("{pre}weight.npy"), self.weight.mut_data())?; - npz_fread(r, format!("{pre}bias.npy"), self.bias.mut_data())?; + fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + npz_fread(r, format!("{p}weight.npy"), self.weight.mut_data())?; + npz_fread(r, format!("{p}bias.npy"), self.bias.mut_data())?; Ok(()) } } -#[cfg(feature = "numpy")] impl SaveToNpz for GeneralizedResidual { - /// Pass through to `F`/`R`'s [SaveToNpz]. fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { self.f.write(&format!("{p}.f"), w)?; self.r.write(&format!("{p}.r"), w) } } -#[cfg(feature = "numpy")] impl LoadFromNpz for GeneralizedResidual { - /// Pass through to `F`/`R`'s [LoadFromNpz]. fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { self.f.read(&format!("{p}.f"), r)?; self.r.read(&format!("{p}.r"), r) } } -#[cfg(feature = "numpy")] impl SaveToNpz for LayerNorm1D { - /// Saves [Self::gamma] to `{pre}gamma.npy` and [Self::beta] to `{pre}beta.npy` - /// using [npz_fwrite()]. - fn write(&self, pre: &str, w: &mut zip::ZipWriter) -> ZipResult<()> { - npz_fwrite(w, format!("{pre}gamma.npy"), self.gamma.data())?; - npz_fwrite(w, format!("{pre}beta.npy"), self.beta.data())?; + fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { + npz_fwrite(w, format!("{p}gamma.npy"), self.gamma.data())?; + npz_fwrite(w, format!("{p}beta.npy"), self.beta.data())?; Ok(()) } } -#[cfg(feature = "numpy")] impl LoadFromNpz for LayerNorm1D { - /// Reads [Self::gamma] from `{p}gamma.npy` and [Self::beta] from `{p}beta.npy` - /// using [npz_fread()]. fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { npz_fread(r, format!("{p}gamma.npy"), self.gamma.mut_data())?; npz_fread(r, format!("{p}beta.npy"), self.beta.mut_data())?; @@ -92,27 +78,17 @@ impl LoadFromNpz for LayerNorm1D { } impl SaveToNpz for Linear { - /// Saves [Self::weight] to `{pre}weight.npy` and [Self::bias] to `{pre}bias.npy` - /// using [npz_fwrite()]. - fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> - where - W: Write + Seek, - { - npz_fwrite(w, format!("{pre}weight.npy"), self.weight.data())?; - npz_fwrite(w, format!("{pre}bias.npy"), self.bias.data())?; + fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { + npz_fwrite(w, format!("{p}weight.npy"), self.weight.data())?; + npz_fwrite(w, format!("{p}bias.npy"), self.bias.data())?; Ok(()) } } impl LoadFromNpz for Linear { - /// Reads [Self::weight] from `{pre}weight.npy` and [Self::bias] from `{pre}bias.npy` - /// using [npz_fread()]. - fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> - where - R: Read + Seek, - { - npz_fread(r, format!("{pre}weight.npy"), self.weight.mut_data())?; - npz_fread(r, format!("{pre}bias.npy"), self.bias.mut_data())?; + fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + npz_fread(r, format!("{p}weight.npy"), self.weight.mut_data())?; + npz_fread(r, format!("{p}bias.npy"), self.bias.mut_data())?; Ok(()) } } @@ -120,11 +96,6 @@ impl LoadFromNpz for Linear { macro_rules! tuple_npz_impl { ([$($name:ident),+], [$($idx:tt),+]) => { impl<$($name: SaveToNpz),+> SaveToNpz for ($($name,)+) { - /// Calls `SaveToNpz::write(self., ...)` on each part of the tuple. See [SaveToNpz]. - /// - /// E.g. for a two tuple (A, B) with `base == ""`, this will call: - /// 1. `self.0.write("0.", w)` - /// 2. `self.1.write("1.", w)` fn write(&self, base: &str, w: &mut ZipWriter) -> ZipResult<()> { $(self.$idx.write(&format!("{}{}.", base, $idx), w)?;)+ Ok(()) @@ -132,11 +103,6 @@ impl<$($name: SaveToNpz),+> SaveToNpz for ($($name,)+) { } impl<$($name: LoadFromNpz),+> LoadFromNpz for ($($name,)+) { - /// Calls `LoadFromNpz::read(self., ...)` on each part of the tuple. See [LoadFromNpz]. - /// - /// E.g. for a two tuple (A, B) with `base == ""`, this will call: - /// 1. `self.0.read("0.", r)` - /// 2. `self.1.read("1.", r)` fn read(&mut self, base: &str, r: &mut ZipArchive) -> Result<(), NpzError> { $(self.$idx.read(&format!("{}{}.", base, $idx), r)?;)+ Ok(()) @@ -151,71 +117,44 @@ tuple_npz_impl!([A, B, C, D], [0, 1, 2, 3]); tuple_npz_impl!([A, B, C, D, E], [0, 1, 2, 3, 4]); tuple_npz_impl!([A, B, C, D, E, F], [0, 1, 2, 3, 4, 5]); -#[cfg(feature = "numpy")] impl SaveToNpz for Repeated { - /// Calls `SaveToNpz::write(self.modules[i], ...)` on each sub module. See [SaveToNpz]. - /// - /// E.g. for a two items with `base == ""`, this will call: - /// 1. `self.modules[0].write("0.", w)` - /// 2. `self.modules[1].write("1.", w)` - fn write(&self, base: &str, w: &mut ZipWriter) -> ZipResult<()> { + fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { for i in 0..N { - self.modules[i].write(&format!("{base}{i}."), w)?; + self.modules[i].write(&format!("{p}{i}."), w)?; } Ok(()) } } -#[cfg(feature = "numpy")] impl LoadFromNpz for Repeated { - /// Calls `LoadFromNpz::read(self.modules[i], ...)` on each sub module. See [LoadFromNpz]. - /// - /// E.g. for a two items with `base == ""`, this will call: - /// 1. `self.modules[0].read("0.", r)` - /// 2. `self.modules[1].read("1.", r)` - fn read(&mut self, base: &str, r: &mut ZipArchive) -> Result<(), NpzError> - where - R: Read + Seek, - { + fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { for i in 0..N { - self.modules[i].read(&format!("{base}{i}."), r)?; + self.modules[i].read(&format!("{p}{i}."), r)?; } Ok(()) } } -#[cfg(feature = "numpy")] impl SaveToNpz for Residual { - /// Pass through to `F`'s [SaveToNpz]. fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { self.0.write(p, w) } } -#[cfg(feature = "numpy")] impl LoadFromNpz for Residual { - /// Pass through to `F`'s [LoadFromNpz]. fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { self.0.read(p, r) } } -#[cfg(feature = "numpy")] impl SaveToNpz for SplitInto { - fn write(&self, p: &str, w: &mut zip::ZipWriter) -> zip::result::ZipResult<()> - where - W: std::io::Write + std::io::Seek, - { + fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { self.0.write(p, w) } } -#[cfg(feature = "numpy")] impl LoadFromNpz for SplitInto { - fn read(&mut self, p: &str, r: &mut zip::ZipArchive) -> Result<(), NpzError> - where - R: std::io::Read + std::io::Seek, - { + fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { self.0.read(p, r) } } @@ -224,8 +163,8 @@ impl LoadFromNpz for SplitInto { impl SaveToNpz for TransformerDecoder { - fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { - self.0.write(pre, w) + fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { + self.0.write(p, w) } } @@ -233,14 +172,14 @@ impl SaveToNpz impl SaveToNpz for TransformerDecoderBlock { - fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { - self.self_attn.write(&format!("{pre}self_attn."), w)?; - self.norm1.write(&format!("{pre}norm1."), w)?; - self.mh_attn.write(&format!("{pre}mh_attn."), w)?; - self.norm2.write(&format!("{pre}norm2."), w)?; - self.ff.0 .0.write(&format!("{pre}linear1."), w)?; - self.ff.0 .2.write(&format!("{pre}linear2."), w)?; - self.norm3.write(&format!("{pre}norm3."), w)?; + fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { + self.self_attn.write(&format!("{p}self_attn."), w)?; + self.norm1.write(&format!("{p}norm1."), w)?; + self.mh_attn.write(&format!("{p}mh_attn."), w)?; + self.norm2.write(&format!("{p}norm2."), w)?; + self.ff.0 .0.write(&format!("{p}linear1."), w)?; + self.ff.0 .2.write(&format!("{p}linear2."), w)?; + self.norm3.write(&format!("{p}norm3."), w)?; Ok(()) } } @@ -274,12 +213,12 @@ impl LoadFromNpz impl SaveToNpz for TransformerEncoderBlock { - fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { - self.self_attn.write(&format!("{pre}self_attn."), w)?; - self.norm1.write(&format!("{pre}norm1."), w)?; - self.norm2.write(&format!("{pre}norm2."), w)?; - self.ff.0 .0.write(&format!("{pre}linear1."), w)?; - self.ff.0 .2.write(&format!("{pre}linear2."), w)?; + fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { + self.self_attn.write(&format!("{p}self_attn."), w)?; + self.norm1.write(&format!("{p}norm1."), w)?; + self.norm2.write(&format!("{p}norm2."), w)?; + self.ff.0 .0.write(&format!("{p}linear1."), w)?; + self.ff.0 .2.write(&format!("{p}linear2."), w)?; Ok(()) } } @@ -288,12 +227,12 @@ impl SaveToNpz impl LoadFromNpz for TransformerEncoderBlock { - fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - self.self_attn.read(&format!("{pre}self_attn."), r)?; - self.norm1.read(&format!("{pre}norm1."), r)?; - self.norm2.read(&format!("{pre}norm2."), r)?; - self.ff.0 .0.read(&format!("{pre}linear1."), r)?; - self.ff.0 .2.read(&format!("{pre}linear2."), r)?; + fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + self.self_attn.read(&format!("{p}self_attn."), r)?; + self.norm1.read(&format!("{p}norm1."), r)?; + self.norm2.read(&format!("{p}norm2."), r)?; + self.ff.0 .0.read(&format!("{p}linear1."), r)?; + self.ff.0 .2.read(&format!("{p}linear2."), r)?; Ok(()) } } @@ -302,14 +241,11 @@ impl LoadFromNpz impl SaveToNpz for MultiHeadAttention { - fn write(&self, pre: &str, w: &mut zip::ZipWriter) -> zip::result::ZipResult<()> - where - W: std::io::Write + std::io::Seek, - { - self.w_q.write(&format!("{pre}w_q."), w)?; - self.w_k.write(&format!("{pre}w_k."), w)?; - self.w_v.write(&format!("{pre}w_v."), w)?; - self.w_o.write(&format!("{pre}w_o."), w)?; + fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { + self.w_q.write(&format!("{p}w_q."), w)?; + self.w_k.write(&format!("{p}w_k."), w)?; + self.w_v.write(&format!("{p}w_v."), w)?; + self.w_o.write(&format!("{p}w_o."), w)?; Ok(()) } } @@ -318,14 +254,11 @@ impl SaveToNpz impl LoadFromNpz for MultiHeadAttention { - fn read(&mut self, pre: &str, r: &mut zip::ZipArchive) -> Result<(), NpzError> - where - R: std::io::Read + std::io::Seek, - { - self.w_q.read(&format!("{pre}w_q."), r)?; - self.w_k.read(&format!("{pre}w_k."), r)?; - self.w_v.read(&format!("{pre}w_v."), r)?; - self.w_o.read(&format!("{pre}w_o."), r)?; + fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + self.w_q.read(&format!("{p}w_q."), r)?; + self.w_k.read(&format!("{p}w_k."), r)?; + self.w_v.read(&format!("{p}w_v."), r)?; + self.w_o.read(&format!("{p}w_o."), r)?; Ok(()) } } @@ -334,9 +267,9 @@ impl LoadFromNpz impl SaveToNpz for Transformer { - fn write(&self, pre: &str, w: &mut ZipWriter) -> ZipResult<()> { - self.encoder.write(&format!("{pre}encoder."), w)?; - self.decoder.write(&format!("{pre}decoder."), w)?; + fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { + self.encoder.write(&format!("{p}encoder."), w)?; + self.decoder.write(&format!("{p}decoder."), w)?; Ok(()) } } @@ -345,9 +278,9 @@ impl LoadFromNpz for Transformer { - fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - self.encoder.read(&format!("{pre}encoder."), r)?; - self.decoder.read(&format!("{pre}decoder."), r)?; + fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + self.encoder.read(&format!("{p}encoder."), r)?; + self.decoder.read(&format!("{p}decoder."), r)?; Ok(()) } } From 5fddb31e2580743d9643533a73da556b2dd8cca3 Mon Sep 17 00:00:00 2001 From: Corey Lowman Date: Sat, 15 Oct 2022 12:34:38 -0400 Subject: [PATCH 5/6] Rewrite npz tests --- src/nn/npz_impls.rs | 343 +++++++++----------------------------------- 1 file changed, 69 insertions(+), 274 deletions(-) diff --git a/src/nn/npz_impls.rs b/src/nn/npz_impls.rs index 41995dd2d..82905ea0d 100644 --- a/src/nn/npz_impls.rs +++ b/src/nn/npz_impls.rs @@ -331,337 +331,132 @@ impl LoadFromNpz for MinPool2D + TensorCreator + Clone, + M: Default + ResetParams + Module + SaveToNpz + LoadFromNpz, + >() + where + M::Output: HasArrayData, + ::Array: std::fmt::Debug + PartialEq, + { let mut rng = thread_rng(); - let mut bn: BatchNorm2D<3> = Default::default(); + let x: I = TensorCreator::randn(&mut rng); + let file = NamedTempFile::new().expect("failed to create tempfile"); - assert_eq!(bn.running_mean.data(), &[0.0; 3]); - assert_eq!(bn.running_var.data(), &[1.0; 3]); - assert_eq!(bn.scale.data(), &[1.0; 3]); - assert_eq!(bn.bias.data(), &[0.0; 3]); + let mut saved: M = Default::default(); + let mut loaded: M = Default::default(); - let x1: Tensor3D<3, 4, 5> = TensorCreator::randn(&mut rng); - let g = backward(bn.forward_mut(x1.trace()).exp().mean()); - bn.update(&mut SimpleGradients(g), &mut Default::default()); + saved.reset_params(&mut rng); + let y = saved.forward(x.clone()); - assert_ne!(bn.running_mean.data(), &[0.0; 3]); - assert_ne!(bn.running_var.data(), &[1.0; 3]); - assert_ne!(bn.scale.data(), &[1.0; 3]); - assert_ne!(bn.bias.data(), &[0.0; 3]); + assert_ne!(loaded.forward(x.clone()).data(), y.data()); - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(bn.save(file.path().to_str().unwrap()).is_ok()); + saved.save(file.path()).expect(""); + loaded.load(file.path()).expect(""); - let mut loaded: BatchNorm2D<3> = Default::default(); - assert!(loaded.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded.scale.data(), bn.scale.data()); - assert_eq!(loaded.bias.data(), bn.bias.data()); - assert_eq!(loaded.running_mean.data(), bn.running_mean.data()); - assert_eq!(loaded.running_var.data(), bn.running_var.data()); + assert_eq!(loaded.forward(x).data(), y.data()); } - #[cfg(feature = "nightly")] #[test] - fn test_save_conv2d() { - let model: Conv2D<2, 4, 3> = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let mut zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - { - let weight_file = zip - .by_name("weight.npy") - .expect("failed to find weight.npy file"); - assert!(weight_file.size() > 0); - } - { - let bias_file = zip - .by_name("bias.npy") - .expect("failed to find bias.npy file"); - assert!(bias_file.size() > 0); - } - } - - #[cfg(feature = "nightly")] - #[test] - fn test_load_conv() { + fn test_batchnorm2d_save_load() { let mut rng = thread_rng(); - let mut saved_model: Conv2D<2, 4, 3> = Default::default(); - saved_model.reset_params(&mut rng); - + let x: Tensor3D<3, 4, 5> = TensorCreator::randn(&mut rng); let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded_model: Conv2D<2, 4, 3> = Default::default(); - assert!(loaded_model.weight.data() != saved_model.weight.data()); - assert!(loaded_model.bias.data() != saved_model.bias.data()); - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.weight.data(), saved_model.weight.data()); - assert_eq!(loaded_model.bias.data(), saved_model.bias.data()); - } + let mut saved: BatchNorm2D<3> = Default::default(); + let mut loaded: BatchNorm2D<3> = Default::default(); - #[test] - fn test_save_load_generalized_residual() { - let mut rng = thread_rng(); - let mut saved_model: GeneralizedResidual, Linear<5, 3>> = Default::default(); - saved_model.reset_params(&mut rng); + saved.running_mean.randomize(&mut rng, &Standard); + saved.running_var.randomize(&mut rng, &Standard); + saved.scale.randomize(&mut rng, &Standard); + saved.bias.randomize(&mut rng, &Standard); + let y = saved.forward(x.clone()); - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + assert_ne!(loaded.forward(x.clone()).data(), y.data()); - let mut loaded_model: GeneralizedResidual, Linear<5, 3>> = Default::default(); - assert_ne!(loaded_model.f.weight.data(), saved_model.f.weight.data()); - assert_ne!(loaded_model.f.bias.data(), saved_model.f.bias.data()); - assert_ne!(loaded_model.r.weight.data(), saved_model.r.weight.data()); - assert_ne!(loaded_model.r.bias.data(), saved_model.r.bias.data()); + saved.save(file.path()).expect(""); + loaded.load(file.path()).expect(""); - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.f.weight.data(), saved_model.f.weight.data()); - assert_eq!(loaded_model.f.bias.data(), saved_model.f.bias.data()); - assert_eq!(loaded_model.r.weight.data(), saved_model.r.weight.data()); - assert_eq!(loaded_model.r.bias.data(), saved_model.r.bias.data()); + assert_eq!(loaded.forward(x).data(), y.data()); } + #[cfg(feature = "nightly")] #[test] - fn test_save_linear() { - let model: Linear<5, 3> = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let mut zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - { - let weight_file = zip - .by_name("weight.npy") - .expect("failed to find weight.npy file"); - assert!(weight_file.size() > 0); - } - { - let bias_file = zip - .by_name("bias.npy") - .expect("failed to find bias.npy file"); - assert!(bias_file.size() > 0); - } + fn test_save_load_conv() { + type T = Conv2D<2, 4, 3>; + test_save_load::, T>(); } #[test] - fn test_load_linear() { - let mut rng = thread_rng(); - let mut saved_model: Linear<5, 3> = Default::default(); - saved_model.reset_params(&mut rng); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded_model: Linear<5, 3> = Default::default(); - assert!(loaded_model.weight.data() != saved_model.weight.data()); - assert!(loaded_model.bias.data() != saved_model.bias.data()); - - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.weight.data(), saved_model.weight.data()); - assert_eq!(loaded_model.bias.data(), saved_model.bias.data()); + fn test_save_load_generalized_residual() { + type T = GeneralizedResidual, Linear<5, 5>>; + test_save_load::, T>(); + test_save_load::, (T, T)>(); } #[test] - fn test_save_tuple() { - let model: ( - Linear<1, 2>, - ReLU, - Linear<2, 3>, - (Dropout, Linear<1, 2>, Linear<3, 4>), - ) = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - let mut names = zip.file_names().collect::>(); - names.sort_unstable(); - assert_eq!( - &names, - &[ - "0.bias.npy", - "0.weight.npy", - "2.bias.npy", - "2.weight.npy", - "3.1.bias.npy", - "3.1.weight.npy", - "3.2.bias.npy", - "3.2.weight.npy", - ] - ); + fn test_save_load_linear() { + type T = Linear<5, 5>; + test_save_load::, T>(); + test_save_load::, (T, T)>(); } #[test] - fn test_load_tuple() { + fn test_save_load_tuple() { type Model = ( Linear<1, 2>, ReLU, Linear<2, 3>, - (Dropout, Linear<1, 2>, Linear<3, 4>), - ); - - let mut rng = thread_rng(); - let mut saved_model: Model = Default::default(); - saved_model.reset_params(&mut rng); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded_model: Model = Default::default(); - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.0.weight.data(), saved_model.0.weight.data()); - assert_eq!(loaded_model.0.bias.data(), saved_model.0.bias.data()); - assert_eq!(loaded_model.2.weight.data(), saved_model.2.weight.data()); - assert_eq!(loaded_model.2.bias.data(), saved_model.2.bias.data()); - assert_eq!( - loaded_model.3 .1.weight.data(), - saved_model.3 .1.weight.data() + (Dropout, Linear<3, 3>, Linear<3, 4>), ); - assert_eq!(loaded_model.3 .1.bias.data(), saved_model.3 .1.bias.data()); - - assert_eq!( - loaded_model.3 .2.weight.data(), - saved_model.3 .2.weight.data() - ); - assert_eq!(loaded_model.3 .2.bias.data(), saved_model.3 .2.bias.data()); + test_save_load::, Model>(); } #[test] - fn test_save_layer_norm() { - let model: LayerNorm1D<13> = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - let mut names = zip.file_names().collect::>(); - names.sort_unstable(); - assert_eq!(&names, &["beta.npy", "gamma.npy",]); - } + fn test_save_load_layer_norm() { + type M = LayerNorm1D<3>; - #[test] - fn test_save_layer_norm_tuple() { - let model: (LayerNorm1D<5>, LayerNorm1D<13>) = Default::default(); + let mut rng = thread_rng(); + let x: Tensor1D<3> = TensorCreator::randn(&mut rng); let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - let mut names = zip.file_names().collect::>(); - names.sort_unstable(); - assert_eq!( - &names, - &["0.beta.npy", "0.gamma.npy", "1.beta.npy", "1.gamma.npy"] - ); - } - #[test] - fn test_load_layer_norm() { - let mut rng = thread_rng(); - let mut saved_model: LayerNorm1D<13> = Default::default(); - saved_model.gamma.randomize(&mut rng, &Standard); - saved_model.beta.randomize(&mut rng, &Standard); + let mut saved: M = Default::default(); + let mut loaded: M = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); + saved.gamma.randomize(&mut rng, &Standard); + saved.beta.randomize(&mut rng, &Standard); + let y = saved.forward(x.clone()); - let mut loaded_model: LayerNorm1D<13> = Default::default(); - assert!(loaded_model.gamma.data() != saved_model.gamma.data()); - assert!(loaded_model.beta.data() != saved_model.beta.data()); + assert_ne!(loaded.forward(x.clone()).data(), y.data()); - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.gamma.data(), saved_model.gamma.data()); - assert_eq!(loaded_model.beta.data(), saved_model.beta.data()); - } + saved.save(file.path()).expect(""); + loaded.load(file.path()).expect(""); - #[test] - fn test_save_repeated() { - let model: Repeated, 4> = Default::default(); - let file = NamedTempFile::new().expect("failed to create tempfile"); - model - .save(file.path().to_str().unwrap()) - .expect("failed to save model"); - let f = File::open(file.path()).expect("failed to open resulting file"); - let zip = ZipArchive::new(f).expect("failed to create zip archive from file"); - let mut names = zip.file_names().collect::>(); - names.sort_unstable(); - assert_eq!( - &names, - &[ - "0.bias.npy", - "0.weight.npy", - "1.bias.npy", - "1.weight.npy", - "2.bias.npy", - "2.weight.npy", - "3.bias.npy", - "3.weight.npy", - ] - ); + assert_eq!(loaded.forward(x).data(), y.data()); } #[test] - fn test_load_repeated() { - type Model = Repeated, 4>; - - let mut rng = thread_rng(); - let mut saved_model: Model = Default::default(); - saved_model.reset_params(&mut rng); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded_model: Model = Default::default(); - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - for i in 0..4 { - assert_eq!( - loaded_model.modules[i].weight.data(), - saved_model.modules[i].weight.data() - ); - assert_eq!( - loaded_model.modules[i].bias.data(), - saved_model.modules[i].bias.data() - ); - } + fn test_save_load_repeated() { + type T = Repeated, 4>; + test_save_load::, T>(); + test_save_load::, (T, T)>(); } #[test] fn test_save_load_residual() { - let mut rng = thread_rng(); - let mut saved_model: Residual> = Default::default(); - saved_model.reset_params(&mut rng); - - let file = NamedTempFile::new().expect("failed to create tempfile"); - assert!(saved_model.save(file.path().to_str().unwrap()).is_ok()); - - let mut loaded_model: Residual> = Default::default(); - assert_ne!(loaded_model.0.weight.data(), saved_model.0.weight.data()); - assert_ne!(loaded_model.0.bias.data(), saved_model.0.bias.data()); - - assert!(loaded_model.load(file.path().to_str().unwrap()).is_ok()); - assert_eq!(loaded_model.0.weight.data(), saved_model.0.weight.data()); - assert_eq!(loaded_model.0.bias.data(), saved_model.0.bias.data()); + type T = Residual>; + test_save_load::, T>(); + test_save_load::, (T, T)>(); } #[cfg(feature = "nightly")] #[test] - fn test_save_and_load() { + fn test_save_load_mha() { let mut rng = thread_rng(); let mut saved: MultiHeadAttention<12, 4> = Default::default(); @@ -684,7 +479,7 @@ mod tests { #[cfg(feature = "nightly")] #[test] - fn test_save_load() { + fn test_save_load_transformer() { let mut rng = thread_rng(); let mut saved: Transformer<16, 4, 3, 4, 8> = Default::default(); From 8e69ba4d7c4b73d812665ecddef29ba307252af0 Mon Sep 17 00:00:00 2001 From: Corey Lowman Date: Sat, 15 Oct 2022 12:37:43 -0400 Subject: [PATCH 6/6] Changing passthrough npz impls to write to .0 --- src/nn/npz_impls.rs | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/nn/npz_impls.rs b/src/nn/npz_impls.rs index 82905ea0d..f3de13486 100644 --- a/src/nn/npz_impls.rs +++ b/src/nn/npz_impls.rs @@ -96,15 +96,15 @@ impl LoadFromNpz for Linear { macro_rules! tuple_npz_impl { ([$($name:ident),+], [$($idx:tt),+]) => { impl<$($name: SaveToNpz),+> SaveToNpz for ($($name,)+) { - fn write(&self, base: &str, w: &mut ZipWriter) -> ZipResult<()> { - $(self.$idx.write(&format!("{}{}.", base, $idx), w)?;)+ + fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { + $(self.$idx.write(&format!("{p}{}.", $idx), w)?;)+ Ok(()) } } impl<$($name: LoadFromNpz),+> LoadFromNpz for ($($name,)+) { - fn read(&mut self, base: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - $(self.$idx.read(&format!("{}{}.", base, $idx), r)?;)+ + fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + $(self.$idx.read(&format!("{p}{}.", $idx), r)?;)+ Ok(()) } } @@ -137,25 +137,25 @@ impl LoadFromNpz for Repeated { impl SaveToNpz for Residual { fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { - self.0.write(p, w) + self.0.write(&format!("{p}.0"), w) } } impl LoadFromNpz for Residual { fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - self.0.read(p, r) + self.0.read(&format!("{p}.0"), r) } } impl SaveToNpz for SplitInto { fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { - self.0.write(p, w) + self.0.write(&format!("{p}.0"), w) } } impl LoadFromNpz for SplitInto { fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - self.0.read(p, r) + self.0.read(&format!("{p}.0"), r) } } @@ -164,7 +164,7 @@ impl SaveToNpz for TransformerDecoder { fn write(&self, p: &str, w: &mut ZipWriter) -> ZipResult<()> { - self.0.write(p, w) + self.0.write(&format!("{p}.0"), w) } } @@ -204,8 +204,8 @@ impl LoadFromNpz impl LoadFromNpz for TransformerDecoder { - fn read(&mut self, pre: &str, r: &mut ZipArchive) -> Result<(), NpzError> { - self.0.read(pre, r) + fn read(&mut self, p: &str, r: &mut ZipArchive) -> Result<(), NpzError> { + self.0.read(&format!("{p}.0"), r) } }