dnrops
/
burn
mirror of https://github.com/tracel-ai/burn.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactor/backend ndarray (#105)

This commit is contained in:
Nathaniel Simard 2022-11-18 20:37:38 -05:00 committed by GitHub
parent 713f078602
commit d45d674a04
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
39 changed files with 426 additions and 385 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
.github/workflows/test-burn-ndarray.yml vendored Normal file
View File

@ -0,0 +1,38 @@
name: test
on: [push]
jobs:
publish:
name: test burn ndarray
runs-on: ubuntu-latest
steps:
- name: checkout
uses: actions/checkout@v2
- name: install rust nightly
uses: actions-rs/toolchain@v1
with:
profile: minimal
toolchain: stable
components: rustfmt, clippy
override: true
- name: check format
run: |
cd burn-ndarray
cargo fmt --check --all
- name: check doc
run: |
cd burn-ndarray
cargo test --doc
- name: check tests
run: |
cd burn-ndarray
cargo test --tests
- name: check clippy
run: |
cargo clippy -p burn-ndarray -- -D warnings

8
.github/workflows/test-burn-tensor.yml vendored
View File

@ -28,11 +28,11 @@ jobs:
cd burn-tensor
cargo test --doc
- name: check tests backend ndarray
- name: check tests
run: |
cd burn-tensor
cargo test --no-default-features --features ndarray --tests
cargo test --tests
- name: check clippy backend ndarray
- name: check clippy backend
run: |
cargo clippy -p burn-tensor --no-default-features --features ndarray -- -D warnings
cargo clippy -p burn-tensor -- -D warnings

1
Cargo.toml
View File

@ -6,5 +6,6 @@ members = [
"burn-tensor-testgen",
"burn-dataset",
"burn-tch",
"burn-ndarray",
"examples/*",
]

19
README.md
View File

@ -54,13 +54,18 @@ The [MNIST](https://github.com/burn-rs/burn/blob/main/examples/mnist) example is
The example can be run like so:
```console
$ git clone https://github.com/burn-rs/burn.git
$ cd burn
$ export TORCH_CUDA_VERSION=cu113 # Set the cuda version
$ # Use the --release flag to really speed up training.
$ cargo run --example mnist --release # CPU NdArray Backend
$ cargo run --example mnist_cuda_gpu --release # GPU Tch Backend
```bash
git clone https://github.com/burn-rs/burn.git
cd burn
# Use the --release flag to really speed up training.
echo "Using ndarray backend"
cargo run --example mnist --release --features ndarray # CPU NdArray Backend - f32 - single thread
cargo run --example mnist --release --features ndarray-blas-openblas # CPU NdArray Backend - f32 - blas with openblas
cargo run --example mnist --release --features ndarray-blas-netlib # CPU NdArray Backend - f32 - blas with netlib
echo "Using tch backend"
export TORCH_CUDA_VERSION=cu113 # Set the cuda version
cargo run --example mnist --release --features tch-gpu # GPU Tch Backend - f16
cargo run --example mnist --release --features tch-cpu # CPU Tch Backend - f32
```
### Components

32
burn-ndarray/Cargo.toml Normal file
View File

@ -0,0 +1,32 @@
[package]
name = "burn-ndarray"
version = "0.2.3"
authors = ["nathanielsimard <nathaniel.simard.42@gmail.com>"]
description = "NdArray backend for burn"
repository = "https://github.com/burn-rs/burn/tree/main/burn-ndarray"
readme="README.md"
keywords = ["deep-learning", "machine-learning", "data"]
categories = ["science"]
license = "MIT/Apache-2.0"
edition = "2021"
[features]
default = []
blas-netlib = ["ndarray/blas", "blas-src/netlib"]
blas-openblas = ["ndarray/blas", "blas-src/openblas", "openblas-src"]
blas-openblas-system = ["ndarray/blas", "blas-src/openblas", "openblas-src/system"]
[dependencies]
burn-tensor = { path = "../burn-tensor", version = "0.2.3" }
blas-src = { version = "0.8.0", default-features = false, optional = true }
openblas-src = { version = "0.10", optional = true }
ndarray = "0.15"
libm = "0.2"
derive-new = "0.5"
rand = "0.8"
num-traits = "0.2"
[dev-dependencies]
burn-tensor = { path = "../burn-tensor", version = "0.2.3", features = ["export_tests"] }

1
burn-ndarray/LICENSE-APACHE Symbolic link
View File

@ -0,0 +1 @@
../LICENSE-APACHE

1
burn-ndarray/LICENSE-MIT Symbolic link
View File

@ -0,0 +1 @@
../LICENSE-MIT

3
burn-ndarray/README.md Normal file
View File

@ -0,0 +1,3 @@
# Burn-NdArray
NdArray backend for burn.

6
burn-tensor/src/tensor/backend/ndarray/backend.rs → burn-ndarray/src/backend.rs
View File

@ -1,8 +1,8 @@
use super::element::NdArrayElement;
use super::NdArrayTensor;
use crate::tensor::backend::Backend;
use crate::tensor::Data;
use crate::{Distribution, Shape};
use burn_tensor::backend::Backend;
use burn_tensor::Data;
use burn_tensor::{Distribution, Shape};
use rand::rngs::StdRng;
use rand::SeedableRng;
use std::sync::Mutex;

2
burn-tensor/src/tensor/backend/ndarray/element.rs → burn-ndarray/src/element.rs
View File

@ -1,4 +1,4 @@
use crate::Element;
use burn_tensor::Element;
pub(crate) trait NdArrayElement:
Element + ndarray::LinalgScalar + ndarray::ScalarOperand + ExpElement + num_traits::FromPrimitive

26
burn-ndarray/src/lib.rs Normal file
View File

@ -0,0 +1,26 @@
#[macro_use]
extern crate derive_new;
#[cfg(any(
feature = "blas-netlib",
feature = "blas-openblas",
feature = "blas-openblas-system",
))]
extern crate blas_src;
mod backend;
mod element;
mod module_ops;
mod ops;
mod tensor;
mod tensor_ops;
pub use backend::*;
pub(crate) use tensor::*;
#[cfg(test)]
mod tests {
type TestBackend = crate::NdArrayBackend<f32>;
burn_tensor::testgen_all!();
}

2
burn-tensor/src/tensor/backend/ndarray/module_ops.rs → burn-ndarray/src/module_ops.rs
View File

@ -1,5 +1,5 @@
use super::{element::NdArrayElement, NdArrayBackend, NdArrayTensor};
use crate::{ops::*, Shape};
use burn_tensor::{ops::*, Shape};
use std::ops::Add;
impl<E: NdArrayElement> ModuleOps<NdArrayBackend<E>> for NdArrayBackend<E> {

3
burn-tensor/src/tensor/backend/ndarray/ops/creation.rs → burn-ndarray/src/ops/creation.rs
View File

@ -1,4 +1,5 @@
use crate::tensor::{backend::ndarray::NdArrayTensor, ops::*, Data};
use crate::NdArrayTensor;
use burn_tensor::{ops::*, Data};
impl<P, const D: usize> Zeros for NdArrayTensor<P, D>
where

0
burn-tensor/src/tensor/backend/ndarray/ops/mod.rs → burn-ndarray/src/ops/mod.rs
View File

34
burn-tensor/src/tensor/backend/ndarray/tensor.rs → burn-ndarray/src/tensor.rs
View File

@ -1,8 +1,5 @@
use super::{element::NdArrayElement, NdArrayBackend};
use crate::{
ops::TensorOps,
tensor::{Data, Shape},
};
use burn_tensor::{ops::TensorOps, Data, Shape};
use ndarray::{s, ArcArray, Array, Axis, Dim, Ix2, Ix3, IxDyn};
#[derive(Debug, Clone)]
@ -22,7 +19,7 @@ impl<E: NdArrayElement, const D: usize> std::ops::Add for NdArrayTensor<E, D> {
#[cfg(test)]
mod utils {
use super::*;
use crate::{backend::NdArrayBackend, ops::TensorOps};
use crate::NdArrayBackend;
impl<E, const D: usize> NdArrayTensor<E, D>
where
@ -38,14 +35,14 @@ mod utils {
}
#[derive(new)]
pub struct BatchMatrix<E, const D: usize> {
pub(crate) struct BatchMatrix<E, const D: usize> {
pub arrays: Vec<ArcArray<E, Ix2>>,
pub shape: Shape<D>,
}
impl<E, const D: usize> BatchMatrix<E, D>
where
E: Clone,
E: NdArrayElement,
{
pub fn from_ndarray(array: ArcArray<E, IxDyn>, shape: Shape<D>) -> Self {
let mut arrays = Vec::new();
@ -66,6 +63,22 @@ where
Self { arrays, shape }
}
pub fn matmul(self, other: BatchMatrix<E, D>) -> Self {
let self_iter = self.arrays.iter();
let other_iter = other.arrays.iter();
let arrays = self_iter
.zip(other_iter)
.map(|(lhs, rhs)| lhs.dot(rhs))
.map(|output| output.into_shared())
.collect();
let mut shape = self.shape;
shape.dims[D - 1] = other.shape.dims[D - 1];
Self::new(arrays, shape)
}
}
fn batch_size<const D: usize>(shape: &Shape<D>) -> usize {
@ -115,7 +128,7 @@ impl<E, const D: usize> NdArrayTensor<E, D>
where
E: Default + Clone,
{
pub fn from_bmatrix(bmatrix: BatchMatrix<E, D>) -> NdArrayTensor<E, D> {
pub(crate) fn from_bmatrix(bmatrix: BatchMatrix<E, D>) -> NdArrayTensor<E, D> {
let shape = bmatrix.shape;
let to_array = |data: BatchMatrix<E, D>| {
let dims = data.shape.dims;
@ -161,10 +174,9 @@ where
#[cfg(test)]
mod tests {
use rand::{rngs::StdRng, SeedableRng};
use super::*;
use crate::tensor::Distribution;
use burn_tensor::Distribution;
use rand::{rngs::StdRng, SeedableRng};
#[test]
fn should_support_into_and_from_data_1d() {

20
burn-tensor/src/tensor/backend/ndarray/tensor_ops.rs → burn-ndarray/src/tensor_ops.rs
View File

@ -1,9 +1,6 @@
use super::{element::NdArrayElement, BatchMatrix, NdArrayBackend, NdArrayTensor};
use crate::{
backend::{Backend, NdArrayDevice},
ops::TensorOps,
to_nd_array_tensor, Data, ElementConversion, Shape,
};
use crate::{to_nd_array_tensor, NdArrayDevice};
use burn_tensor::{backend::Backend, ops::TensorOps, Data, ElementConversion, Shape};
use ndarray::{Axis, Dim, IxDyn, SliceInfoElem};
use std::{cmp::Ordering, ops::Range};
@ -184,18 +181,7 @@ impl<E: NdArrayElement> TensorOps<NdArrayBackend<E>> for NdArrayBackend<E> {
) -> <NdArrayBackend<E> as Backend>::TensorPrimitive<D> {
let batch_self = BatchMatrix::from_ndarray(lhs.array.clone(), lhs.shape);
let batch_other = BatchMatrix::from_ndarray(rhs.array.clone(), rhs.shape);
let self_iter = batch_self.arrays.iter();
let other_iter = batch_other.arrays.iter();
let arrays = self_iter
.zip(other_iter)
.map(|(lhs, rhs)| lhs.dot(rhs))
.map(|output| output.into_shared())
.collect();
let mut shape = lhs.shape;
shape.dims[D - 1] = rhs.shape.dims[D - 1];
let output = BatchMatrix::new(arrays, shape);
let output = batch_self.matmul(batch_other);
NdArrayTensor::from_bmatrix(output)
}

2
burn-tch/Cargo.toml
View File

@ -17,9 +17,7 @@ doc = ["tch/doc-only"]
[dependencies]
burn-tensor = { path = "../burn-tensor", version = "0.2.3", default-features = false }
rand = "0.8"
num-traits = "0.2"
tch = { version = "0.8" }
serde = { version = "1.0", features = ["derive"] }
lazy_static = "1.4"
half = { version = "1.6", features = ["num-traits"] } # needs to be 1.6 to work with tch

21
burn-tensor-testgen/src/lib.rs
View File

@ -23,24 +23,5 @@ pub fn testgen(attr: TokenStream, item: TokenStream) -> TokenStream {
}
};
let test_gen = quote! {
#[cfg(test)]
use crate::tests::TestBackend;
#[cfg(test)]
use crate as burn_tensor;
#[cfg(test)]
type TestADBackend = burn_tensor::backend::ADBackendDecorator<TestBackend>;
#[cfg(test)]
type TestTensor<const D: usize> = burn_tensor::Tensor<TestBackend, D>;
#[cfg(test)]
type TestADTensor<const D: usize> = burn_tensor::Tensor<TestADBackend, D>;
#[cfg(test)]
#item
};
quote! {
#test_gen
#macro_gen
}
.into()
macro_gen.into()
}

14
burn-tensor/Cargo.toml
View File

@ -14,15 +14,9 @@ categories = ["science"]
license = "MIT/Apache-2.0"
edition = "2021"
[package.metadata.docs.rs]
features = ["doc"]
all-features = false
no-default-features = true
[features]
default = ["ndarray"]
ndarray = ["dep:ndarray", "dep:libm"]
export_tests = ["dep:burn-tensor-testgen"]
default = []
export_tests = ["burn-tensor-testgen"]
[dependencies]
burn-tensor-testgen = { path = "../burn-tensor-testgen", optional = true }
@ -31,10 +25,6 @@ derive-new = "0.5"
rand = "0.8"
half = { version = "1.6", features = ["num-traits"] } # needs to be 1.6 to work with tch
# NdArray
ndarray = { version = "0.15", optional = true }
libm = { version = "0.2", optional = true }
# Autodiff
nanoid = "0.4"

54
burn-tensor/examples/autodiff_simple.rs
View File

@ -1,54 +0,0 @@
use burn_tensor::{activation, backend, Data, Distribution, Shape, Tensor};
use rand::{rngs::StdRng, SeedableRng};
fn loss<B: backend::Backend>(x: &Tensor<B, 2>, y: &Tensor<B, 2>) -> Tensor<B, 2> {
let z = x.matmul(y);
let z = activation::relu(&z);
println!("fn name : loss");
println!("backend : {}", B::name());
println!("z : {}", z.to_data());
println!("autodiff : {}", B::ad_enabled());
z
}
fn run_ad<B: backend::ADBackend>(x: Data<B::Elem, 2>, y: Data<B::Elem, 2>) {
println!("---------- Ad Enabled -----------");
let x: Tensor<B, 2> = Tensor::from_data(x);
let y: Tensor<B, 2> = Tensor::from_data(y);
let z = loss(&x, &y);
let grads = z.backward();
println!("x_grad : {}", x.grad(&grads).unwrap().to_data());
println!("y_grad : {}", y.grad(&grads).unwrap().to_data());
println!("---------------------------------");
println!()
}
fn run<B: backend::Backend>(x: Data<B::Elem, 2>, y: Data<B::Elem, 2>) {
println!("---------- Ad Disabled ----------");
loss::<B>(&Tensor::from_data(x), &Tensor::from_data(y));
println!("---------------------------------");
println!()
}
fn main() {
// Same data for all backends
let mut rng = StdRng::from_entropy();
let x = Data::random(Shape::new([2, 3]), Distribution::Standard, &mut rng);
let y = Data::random(Shape::new([3, 1]), Distribution::Standard, &mut rng);
#[cfg(feature = "ndarray")]
{
run::<backend::NdArrayBackend<f32>>(x.clone(), y.clone());
run_ad::<backend::NdArrayADBackend<f32>>(x, y);
}
#[cfg(feature = "tch")]
{
run::<backend::TchBackend<f32>>(x.clone(), y.clone());
run_ad::<backend::TchADBackend<f32>>(x, y);
}
}

2
burn-tensor/src/lib.rs
View File

@ -8,8 +8,6 @@ mod tensor;
#[cfg(feature = "export_tests")]
mod tests;
#[cfg(all(test, not(feature = "export_tests")))]
mod tests;
pub use half::f16;
pub use tensor::*;

81
burn-tensor/src/tensor/backend/autodiff/tensor/ad.rs
View File

@ -21,84 +21,3 @@ impl<B: Backend, const D: usize> AsNode<B::TensorPrimitive<D>> for ADTensor<D, B
&self.node
}
}
#[cfg(test)]
mod tests {
use crate::tensor::{backend::autodiff::helper::TestADTensor, Data};
#[test]
fn should_diff_full_complex_1() {
let data_1: Data<f32, 2> = Data::from([[1.0, 7.0], [13.0, -3.0]]);
let data_2: Data<f32, 2> = Data::from([[4.0, 7.0], [2.0, 3.0]]);
let tensor_1 = TestADTensor::from_data(data_1);
let tensor_2 = TestADTensor::from_data(data_2);
let tensor_3 = tensor_1.matmul(&tensor_2);
let tensor_4 = tensor_3.matmul(&tensor_1);
let tensor_5 = tensor_4.mul(&tensor_2);
let grads = tensor_5.backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
assert_eq!(
grad_1.to_data(),
Data::from([[593., 463.0], [487.0, 539.0]])
);
assert_eq!(
grad_2.to_data(),
Data::from([[734.0, 294.0], [1414.0, 242.0]])
);
}
#[test]
fn should_diff_full_complex_2() {
let data_1: Data<f64, 2> = Data::from([[1.0, 7.0], [13.0, -3.0]]);
let data_2: Data<f64, 2> = Data::from([[4.0, 7.0], [2.0, 3.0]]);
let tensor_1 = TestADTensor::from_data(data_1);
let tensor_2 = TestADTensor::from_data(data_2);
let tensor_3 = tensor_1.matmul(&tensor_2);
let tensor_4 = tensor_3.matmul(&tensor_1);
let tensor_5 = tensor_4.add_scalar(17.0).add(&tensor_2);
let grads = tensor_5.backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
assert_eq!(
grad_1.to_data(),
Data::from([[166.0, 110.0], [212.0, 156.0]])
);
assert_eq!(grad_2.to_data(), Data::from([[113.0, 141.0], [33.0, 41.0]]));
}
#[test]
fn should_diff_full_complex_3() {
let data_1: Data<f64, 2> = Data::from([[1.0, 7.0], [13.0, -3.0]]);
let data_2: Data<f64, 2> = Data::from([[4.0, 7.0], [2.0, 3.0]]);
let tensor_1 = TestADTensor::from_data(data_1);
let tensor_2 = TestADTensor::from_data(data_2);
let tensor_3 = tensor_1.matmul(&tensor_2);
let tensor_4 = tensor_3.matmul(&tensor_1);
let tensor_5 = tensor_4.sub(&tensor_2);
let tensor_6 = tensor_5.add(&tensor_4);
let grads = tensor_6.backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
assert_eq!(
grad_1.to_data(),
Data::from([[332.0, 220.0], [424.0, 312.0]])
);
assert_eq!(grad_2.to_data(), Data::from([[223.0, 279.0], [63.0, 79.0]]));
}
}

12
burn-tensor/src/tensor/backend/autodiff/tensor/base.rs
View File

@ -30,15 +30,3 @@ impl<B: Backend, const D: usize> ADTensor<D, B> {
self.node.state.value_ref()
}
}
#[cfg(test)]
pub mod helper {
#[cfg(feature = "ndarray")]
mod helper_impl {
use crate::tensor::backend::autodiff::ADBackendNdArray;
use crate::tensor::Tensor;
pub type TestADTensor<E, const D: usize> = Tensor<ADBackendNdArray<E>, D>;
}
pub use helper_impl::*;
}

1
burn-tensor/src/tensor/backend/autodiff/tensor/mod.rs
View File

@ -1,6 +1,5 @@
mod ad;
mod base;
mod multithread;
pub use ad::*;
pub use base::*;

84
burn-tensor/src/tensor/backend/autodiff/tensor/multithread.rs
View File

@ -1,84 +0,0 @@
#[cfg(test)]
mod tests {
use crate::tensor::{backend::autodiff::helper::TestADTensor, Data};
#[test]
fn should_behave_the_same_with_multithread() {
let data_1: Data<f64, 2> = Data::from([[1.0, 7.0], [13.0, -3.0]]);
let data_2: Data<f64, 2> = Data::from([[4.0, 7.0], [2.0, 3.0]]);
let with_move = || {
let tensor_1 = TestADTensor::from_data(data_1.clone());
let tensor_2 = TestADTensor::from_data(data_2.clone());
let tensor_3 = tensor_1.matmul(&tensor_2);
let tensor_4 = tensor_3.matmul(&tensor_2);
let tensor_5 = tensor_4.matmul(&tensor_3);
// Task 1
let tensor_1_cloned = tensor_1.clone();
let tensor_2_cloned = tensor_2.clone();
let tensor_5_cloned = tensor_5.clone();
let first_call = move || {
let tensor_6_1 = tensor_5_cloned.matmul(&tensor_2_cloned);
tensor_6_1.matmul(&tensor_1_cloned)
};
// Task 2
let tensor_1_cloned = tensor_1.clone();
let tensor_2_cloned = tensor_2.clone();
let tensor_5_cloned = tensor_5;
let second_call = move || {
let tensor_6_2 = tensor_5_cloned.matmul(&tensor_1_cloned);
tensor_6_2.matmul(&tensor_2_cloned)
};
let tensor_7_1_handle = std::thread::spawn(first_call);
let tensor_7_2_handle = std::thread::spawn(second_call);
let tensor_7_1 = tensor_7_1_handle.join().unwrap();
let tensor_7_2 = tensor_7_2_handle.join().unwrap();
let tensor_8 = tensor_7_1.matmul(&tensor_7_2);
let grads = tensor_8.backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
(grad_1, grad_2)
};
let without_move = || {
let tensor_1 = TestADTensor::from_data(data_1.clone());
let tensor_2 = TestADTensor::from_data(data_2.clone());
let tensor_3 = tensor_1.matmul(&tensor_2);
let tensor_4 = tensor_3.matmul(&tensor_2);
let tensor_5 = tensor_4.matmul(&tensor_3);
// Task 1
let tensor_6_1 = tensor_5.matmul(&tensor_2);
let tensor_7_1 = tensor_6_1.matmul(&tensor_1);
// Task 2
let tensor_6_2 = tensor_5.matmul(&tensor_1);
let tensor_7_2 = tensor_6_2.matmul(&tensor_2);
let tensor_8 = tensor_7_1.matmul(&tensor_7_2);
let grads = tensor_8.backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
(grad_1, grad_2)
};
let (grad_1, grad_2) = without_move();
let (grad_1_moved, grad_2_moved) = with_move();
assert_eq!(grad_1.to_data(), grad_1_moved.to_data());
assert_eq!(grad_2.to_data(), grad_2_moved.to_data());
}
}

16
burn-tensor/src/tensor/backend/mod.rs
View File

@ -1,19 +1,7 @@
mod base;
pub use base::*;
pub(crate) mod autodiff;
pub use autodiff::ADBackendDecorator;
// Not needed for now, usefull for different tensor memory layout
// pub mod conversion;
pub(crate) mod autodiff;
#[cfg(feature = "ndarray")]
pub(crate) mod ndarray;
#[cfg(feature = "ndarray")]
pub type NdArrayADBackend<E> = self::autodiff::ADBackendNdArray<E>;
#[cfg(feature = "ndarray")]
pub type NdArrayBackend<E> = self::ndarray::NdArrayBackend<E>;
#[cfg(feature = "ndarray")]
pub type NdArrayDevice = self::ndarray::NdArrayDevice;
pub use autodiff::ADBackendDecorator;

11
burn-tensor/src/tensor/backend/ndarray/mod.rs
View File

@ -1,11 +0,0 @@
mod backend;
mod element;
mod module_ops;
mod ops;
mod shape;
mod tensor;
mod tensor_ops;
pub use backend::*;
pub use shape::*;
pub use tensor::*;

22
burn-tensor/src/tensor/backend/ndarray/shape.rs
View File

@ -1,22 +0,0 @@
use crate::tensor::Shape;
use ndarray::Dim;
macro_rules! define_convertion {
(
$n:expr
) => {
impl From<Shape<$n>> for Dim<[usize; $n]> {
fn from(shape: Shape<$n>) -> Dim<[usize; $n]> {
Dim(shape.dims)
}
}
};
}
define_convertion!(0);
define_convertion!(1);
define_convertion!(2);
define_convertion!(3);
define_convertion!(4);
define_convertion!(5);
define_convertion!(6);

81
burn-tensor/src/tests/grad/complex.rs Normal file
View File

@ -0,0 +1,81 @@
#[burn_tensor_testgen::testgen(ad_complex)]
mod tests {
use super::*;
use burn_tensor::Data;
#[test]
fn should_diff_full_complex_1() {
let data_1: Data<f32, 2> = Data::from([[1.0, 7.0], [13.0, -3.0]]);
let data_2: Data<f32, 2> = Data::from([[4.0, 7.0], [2.0, 3.0]]);
let tensor_1 = TestADTensor::from_data(data_1);
let tensor_2 = TestADTensor::from_data(data_2);
let tensor_3 = tensor_1.matmul(&tensor_2);
let tensor_4 = tensor_3.matmul(&tensor_1);
let tensor_5 = tensor_4.mul(&tensor_2);
let grads = tensor_5.backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
assert_eq!(
grad_1.to_data(),
Data::from([[593., 463.0], [487.0, 539.0]])
);
assert_eq!(
grad_2.to_data(),
Data::from([[734.0, 294.0], [1414.0, 242.0]])
);
}
#[test]
fn should_diff_full_complex_2() {
let data_1: Data<f32, 2> = Data::from([[1.0, 7.0], [13.0, -3.0]]);
let data_2: Data<f32, 2> = Data::from([[4.0, 7.0], [2.0, 3.0]]);
let tensor_1 = TestADTensor::from_data(data_1);
let tensor_2 = TestADTensor::from_data(data_2);
let tensor_3 = tensor_1.matmul(&tensor_2);
let tensor_4 = tensor_3.matmul(&tensor_1);
let tensor_5 = tensor_4.add_scalar(17.0).add(&tensor_2);
let grads = tensor_5.backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
assert_eq!(
grad_1.to_data(),
Data::from([[166.0, 110.0], [212.0, 156.0]])
);
assert_eq!(grad_2.to_data(), Data::from([[113.0, 141.0], [33.0, 41.0]]));
}
#[test]
fn should_diff_full_complex_3() {
let data_1: Data<f32, 2> = Data::from([[1.0, 7.0], [13.0, -3.0]]);
let data_2: Data<f32, 2> = Data::from([[4.0, 7.0], [2.0, 3.0]]);
let tensor_1 = TestADTensor::from_data(data_1);
let tensor_2 = TestADTensor::from_data(data_2);
let tensor_3 = tensor_1.matmul(&tensor_2);
let tensor_4 = tensor_3.matmul(&tensor_1);
let tensor_5 = tensor_4.sub(&tensor_2);
let tensor_6 = tensor_5.add(&tensor_4);
let grads = tensor_6.backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
assert_eq!(
grad_1.to_data(),
Data::from([[332.0, 220.0], [424.0, 312.0]])
);
assert_eq!(grad_2.to_data(), Data::from([[223.0, 279.0], [63.0, 79.0]]));
}
}

2
burn-tensor/src/tests/grad/mod.rs
View File

@ -1,6 +1,7 @@
mod add;
mod aggregation;
mod cat;
mod complex;
mod cross_entropy;
mod div;
mod erf;
@ -10,6 +11,7 @@ mod log;
mod mask;
mod matmul;
mod mul;
mod multithread;
mod neg;
mod pow;
mod relu;

121
burn-tensor/src/tests/grad/multithread.rs Normal file
View File

@ -0,0 +1,121 @@
#[burn_tensor_testgen::testgen(ad_multithread)]
mod tests {
use super::*;
use burn_tensor::Data;
#[test]
fn should_diff_mean() {
let data_1 = Data::<f32, 2>::from([[1.0, 7.0], [-2.0, -3.0]]);
let data_2 = Data::<f32, 2>::from([[4.0, -7.0], [2.0, 3.0]]);
let tensor_1 = TestADTensor::from_data(data_1);
let tensor_2 = TestADTensor::from_data(data_2);
let tensor_3 = tensor_1.matmul(&tensor_2);
let tensor_4 = tensor_1.mul(&tensor_3.mean().unsqueeze());
let grads = tensor_4.backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
grad_1
.to_data()
.assert_approx_eq(&Data::from([[3.5, 9.5], [3.5, 9.5]]), 5);
grad_2
.to_data()
.assert_approx_eq(&Data::from([[-0.75, -0.75], [3.0, 3.0]]), 5);
}
#[test]
fn should_diff_sum_1() {
let data_1 = Data::<f32, 2>::from([[1.0, 7.0], [-2.0, -3.0]]);
let data_2 = Data::<f32, 2>::from([[4.0, -7.0], [2.0, 3.0]]);
let tensor_1 = TestADTensor::from_data(data_1);
let tensor_2 = TestADTensor::from_data(data_2);
let tensor_3 = tensor_1.matmul(&tensor_2);
let tensor_4 = tensor_1.mul(&tensor_3.sum().unsqueeze());
let grads = tensor_4.backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
grad_1
.to_data()
.assert_approx_eq(&Data::from([[14.0, 38.0], [14.0, 38.0]]), 5);
grad_2
.to_data()
.assert_approx_eq(&Data::from([[-3.0, -3.0], [12.0, 12.0]]), 5);
}
#[test]
fn should_diff_sum_2() {
let data_1 = Data::from([[0.0, 1.0], [3.0, 4.0]]);
let data_2 = Data::from([[6.0, 7.0], [9.0, 10.0]]);
let tensor_1 = TestADTensor::from_data(data_1);
let tensor_2 = TestADTensor::from_data(data_2);
let tensor_3 = tensor_1.matmul(&tensor_2);
let tensor_4 = tensor_3.sum_dim(1);
let tensor_5 = tensor_4.mul(&tensor_3);
let grads = tensor_5.sum().backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
grad_1
.to_data()
.assert_approx_eq(&Data::from([[494.0, 722.0], [2990.0, 4370.0]]), 3);
grad_2
.to_data()
.assert_approx_eq(&Data::from([[690.0, 690.0], [958.0, 958.0]]), 3);
}
#[test]
fn should_diff_mean_dim() {
let data_1 = Data::<f32, 2>::from([[1.0, 7.0], [-2.0, -3.0]]);
let data_2 = Data::<f32, 2>::from([[4.0, -7.0], [2.0, 3.0]]);
let tensor_1 = TestADTensor::from_data(data_1);
let tensor_2 = TestADTensor::from_data(data_2);
let tensor_3 = tensor_1.matmul(&tensor_2);
let tensor_4 = tensor_1.mul(&tensor_3.mean_dim(1).unsqueeze());
let grads = tensor_4.backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
grad_1
.to_data()
.assert_approx_eq(&Data::from([[4.0, 36.0], [3.0, -17.0]]), 5);
grad_2
.to_data()
.assert_approx_eq(&Data::from([[9.0, 9.0], [35.5, 35.5]]), 5);
}
#[test]
fn should_diff_sum_dim() {
let data_1 = Data::<f32, 2>::from([[1.0, 7.0], [-2.0, -3.0]]);
let data_2 = Data::<f32, 2>::from([[4.0, -7.0], [2.0, 3.0]]);
let tensor_1 = TestADTensor::from_data(data_1);
let tensor_2 = TestADTensor::from_data(data_2);
let tensor_3 = tensor_1.matmul(&tensor_2);
let tensor_4 = tensor_1.mul(&tensor_3.sum_dim(1).unsqueeze());
let grads = tensor_4.backward();
let grad_1 = tensor_1.grad(&grads).unwrap();
let grad_2 = tensor_2.grad(&grads).unwrap();
grad_1
.to_data()
.assert_approx_eq(&Data::from([[8.0, 72.0], [6.0, -34.0]]), 5);
grad_2
.to_data()
.assert_approx_eq(&Data::from([[18.0, 18.0], [71.0, 71.0]]), 5);
}
}

5
burn-tensor/src/tests/mod.rs
View File

@ -4,9 +4,6 @@ mod module;
mod ops;
mod stats;
#[cfg(test)]
type TestBackend = crate::backend::NdArrayBackend<f32>;
#[macro_export]
macro_rules! testgen_all {
() => {
@ -16,6 +13,8 @@ macro_rules! testgen_all {
burn_tensor::testgen_softmax!();
// test ad
burn_tensor::testgen_ad_complex!();
burn_tensor::testgen_ad_multithread!();
burn_tensor::testgen_ad_add!();
burn_tensor::testgen_ad_aggregation!();
burn_tensor::testgen_ad_cat!();

3
burn/Cargo.toml
View File

@ -11,8 +11,6 @@ license = "MIT/Apache-2.0"
edition = "2021"
[features]
default = ["ndarray"]
ndarray = ["burn-tensor/ndarray"]
[dependencies]
burn-tensor = { path = "../burn-tensor", version = "0.2.3", default-features = false }
@ -45,3 +43,4 @@ nanoid = "0.4"
[dev-dependencies]
burn-dataset = { path = "../burn-dataset", version = "0.2.3", features = ["fake"] }
burn-ndarray = { path = "../burn-ndarray" }

4
burn/src/lib.rs
View File

@ -10,6 +10,6 @@ pub mod tensor;
pub mod train;
#[cfg(test)]
pub type TestBackend = crate::tensor::backend::NdArrayBackend<f32>;
pub type TestBackend = burn_ndarray::NdArrayBackend<f32>;
#[cfg(test)]
pub type TestADBackend = crate::tensor::backend::NdArrayADBackend<f32>;
pub type TestADBackend = crate::tensor::backend::ADBackendDecorator<TestBackend>;

2
burn/tests/derive_module.rs
View File

@ -2,7 +2,7 @@ use burn::module::{Module, Param};
use burn::tensor::backend::Backend;
use burn::tensor::{Distribution, Shape, Tensor};
type TestBackend = burn::tensor::backend::NdArrayBackend<f32>;
pub type TestBackend = burn_ndarray::NdArrayBackend<f32>;
#[derive(Module, Debug)]
struct ModuleBasic<B>

11
examples/mnist/Cargo.toml
View File

@ -6,9 +6,18 @@ license = "MIT/Apache-2.0"
edition = "2021"
publish = false
[features]
default = []
tch-cpu = ["dep:burn-tch"]
tch-gpu = ["dep:burn-tch"]
ndarray = ["burn-ndarray"]
ndarray-blas-netlib = ["burn-ndarray/blas-netlib"]
ndarray-blas-openblas = ["burn-ndarray/blas-openblas"]
[dependencies]
burn = { path = "../../burn" }
burn-tch = { path = "../../burn-tch" }
burn-tch = { path = "../../burn-tch", optional = true }
burn-ndarray = { path = "../../burn-ndarray", optional = true }
# Serialization
serde = { version = "1.0", features = ["derive"] }

55
examples/mnist/examples/mnist.rs
View File

@ -1,9 +1,52 @@
use mnist::training;
#[cfg(any(
feature = "ndarray",
feature = "ndarray-blas-netlib",
feature = "ndarray-blas-openblas",
))]
mod ndarray {
use burn::tensor::backend::ADBackendDecorator;
use burn_ndarray::{NdArrayBackend, NdArrayDevice};
use mnist::training;
pub fn run() {
let device = NdArrayDevice::Cpu;
training::run::<ADBackendDecorator<NdArrayBackend<f32>>>(device);
}
}
#[cfg(feature = "tch-gpu")]
mod tch_gpu {
use burn::tensor::backend::ADBackendDecorator;
use burn_tch::{TchBackend, TchDevice};
use mnist::training;
pub fn run() {
let device = TchDevice::Cuda(0);
training::run::<ADBackendDecorator<TchBackend<burn::tensor::f16>>>(device);
}
}
#[cfg(feature = "tch-cpu")]
mod tch_cpu {
use burn::tensor::backend::ADBackendDecorator;
use burn_tch::{TchBackend, TchDevice};
use mnist::training;
pub fn run() {
let device = TchDevice::Cpu;
training::run::<ADBackendDecorator<TchBackend<f32>>>(device);
}
}
fn main() {
use burn::tensor::backend::{NdArrayADBackend, NdArrayDevice};
let device = NdArrayDevice::Cpu;
training::run::<NdArrayADBackend<f32>>(device);
println!("Done.");
#[cfg(any(
feature = "ndarray",
feature = "ndarray-blas-netlib",
feature = "ndarray-blas-openblas",
))]
ndarray::run();
#[cfg(feature = "tch-gpu")]
tch_gpu::run();
#[cfg(feature = "tch-cpu")]
tch_cpu::run();
}

9
examples/mnist/examples/mnist_cuda_gpu.rs
View File

@ -1,9 +0,0 @@
use burn::tensor::backend::ADBackendDecorator;
use burn_tch::{TchBackend, TchDevice};
use mnist::training;
fn main() {
let device = TchDevice::Cuda(0);
training::run::<ADBackendDecorator<TchBackend<burn::tensor::f16>>>(device);
println!("Done.");
}

2
examples/mnist/src/model.rs
View File

@ -16,7 +16,7 @@ use burn::{
pub struct MnistConfig {
#[config(default = 6)]
pub num_epochs: usize,
#[config(default = 128)]
#[config(default = 64)]
pub batch_size: usize,
#[config(default = 8)]
pub num_workers: usize,