Refactor/extract tch backend (tracel-ai#103)

djdisodo · Nov 16, 2022 · ab51c22 · ab51c22
1 parent 23677b8
commit ab51c22
Show file tree

Hide file tree

Showing 123 changed files with 1,805 additions and 1,637 deletions.
diff --git a/.github/workflows/test-burn-tch.yml b/.github/workflows/test-burn-tch.yml
@@ -0,0 +1,38 @@
+name: test
+
+on: [push]
+
+jobs:
+  publish:
+    name: test burn tch
+    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-tch
+          cargo fmt --check --all
+
+      - name: check doc
+        run: |
+          cd burn-tch
+          cargo test --no-default-features --features doc --doc
+
+      - name: check tests
+        run: |
+          cd burn-tch
+          cargo test --tests
+
+      - name: check clippy
+        run: |
+          cargo clippy -p burn-tch -- -D warnings
diff --git a/.github/workflows/test-burn-tensor.yml b/.github/workflows/test-burn-tensor.yml
@@ -26,22 +26,13 @@ jobs:
       - name: check doc
         run: |
           cd burn-tensor
-          cargo test --no-default-features --features doc --doc
+          cargo test --doc
 
       - name: check tests backend ndarray
         run: |
           cd burn-tensor
-          cargo test --no-default-features --features ndarray --tests
-
-      - name: check tests backend tch
-        run: |
-          cd burn-tensor
-          cargo test --no-default-features --features tch --tests
-
-      - name: check clippy backend tch
-        run: |
-          cargo clippy -p burn-tensor --no-default-features --features tch -- -D warnings
+          cargo test --no-default-features --features ndarray export_tests --tests
 
       - name: check clippy backend ndarray
         run: |
-          cargo clippy -p burn-tensor --no-default-features --features tch -- -D warnings
+          cargo clippy -p burn-tensor --no-default-features --features ndarray -- -D warnings
diff --git a/.github/workflows/test-burn.yml b/.github/workflows/test-burn.yml
@@ -27,7 +27,7 @@ jobs:
       - name: check doc
         run: |
           cd burn
-          cargo test --no-default-features --features doc --doc
+          cargo test --doc
 
       - name: check tests
         run: |

diff --git a/Cargo.toml b/Cargo.toml
@@ -4,5 +4,6 @@ members = [
   "burn-derive",
   "burn-tensor",
   "burn-dataset",
+  "burn-tch",
   "examples/*",
 ]
diff --git a/burn-tch/Cargo.toml b/burn-tch/Cargo.toml
@@ -0,0 +1,27 @@
+[package]
+name = "burn-tch"
+version = "0.2.3"
+authors = ["nathanielsimard <nathaniel.simard.42@gmail.com>"]
+
+description = "Tch backend for burn"
+repository = "https://github.com/burn-rs/burn/tree/main/burn-tch"
+readme="README.md"
+keywords = ["deep-learning", "machine-learning", "data"]
+categories = ["science"]
+license = "MIT/Apache-2.0"
+edition = "2021"
+
+[features]
+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
+
+[dev-dependencies]
+burn-tensor = { path = "../burn-tensor", version = "0.2.3", default-features = false, features = ["export_tests"] }
diff --git a/burn-tch/LICENSE-APACHE b/burn-tch/LICENSE-APACHE
@@ -0,0 +1 @@
+../LICENSE-APACHE
diff --git a/burn-tch/LICENSE-MIT b/burn-tch/LICENSE-MIT
@@ -0,0 +1 @@
+../LICENSE-MIT
diff --git a/burn-tch/README.md b/burn-tch/README.md
@@ -0,0 +1,3 @@
+# Burn-tch
+
+Tch backend for burn.
diff --git a/...-tensor/src/tensor/backend/tch/backend.rs → burn-tch/src/backend.rs b/...-tensor/src/tensor/backend/tch/backend.rs → burn-tch/src/backend.rs
@@ -1,7 +1,7 @@
 use super::element::TchElement;
 use super::TchTensor;
-use crate::tensor::backend::Backend;
-use crate::tensor::{Data, Distribution, Shape};
+use burn_tensor::backend::Backend;
+use burn_tensor::{Data, Distribution, Shape};
 
 #[derive(Clone, Copy, Debug)]
 /// The device struct when using the `tch` backend.
@@ -11,7 +11,7 @@ use crate::tensor::{Data, Distribution, Shape};
 /// # Example
 ///
 /// ```rust
-/// use burn_tensor::backend::TchDevice;
+/// use burn_tch::TchDevice;
 ///
 /// let device_gpu_1 = TchDevice::Cuda(0); // First GPU
 /// let device_gpu_2 = TchDevice::Cuda(1); // Second GPU

diff --git a/burn-tch/src/element.rs b/burn-tch/src/element.rs
@@ -0,0 +1,14 @@
+use burn_tensor::Element;
+use half::f16;
+
+pub trait TchElement: Element + tch::kind::Element {}
+
+impl TchElement for f64 {}
+impl TchElement for f32 {}
+impl TchElement for f16 {}
+
+impl TchElement for i64 {}
+impl TchElement for i32 {}
+impl TchElement for i16 {}
+
+impl TchElement for u8 {}
diff --git a/burn-tensor/src/tensor/backend/tch/mod.rs → burn-tch/src/lib.rs b/burn-tensor/src/tensor/backend/tch/mod.rs → burn-tch/src/lib.rs
@@ -8,3 +8,10 @@ mod tensor_ops;
 pub use backend::*;
 pub use tensor::*;
 pub use tensor_ops::*;
+
+#[cfg(test)]
+mod tests {
+    type TestBackend = crate::TchBackend<f32>;
+
+    burn_tensor::test_all!();
+}
diff --git a/...nsor/src/tensor/backend/tch/module_ops.rs → burn-tch/src/module_ops.rs b/...nsor/src/tensor/backend/tch/module_ops.rs → burn-tch/src/module_ops.rs
@@ -1,5 +1,5 @@
 use super::{element::TchElement, TchBackend, TchTensor};
-use crate::{ops::ModuleOps, Shape};
+use burn_tensor::{ops::ModuleOps, Shape};
 
 impl<E: TchElement> ModuleOps<TchBackend<E>> for TchBackend<E> {
     fn embedding(weights: &TchTensor<E, 2>, indexes: &TchTensor<i64, 2>) -> TchTensor<E, 3> {

diff --git a/...or/src/tensor/backend/tch/ops/creation.rs → burn-tch/src/ops/creation.rs b/...or/src/tensor/backend/tch/ops/creation.rs → burn-tch/src/ops/creation.rs
@@ -1,4 +1,5 @@
-use crate::tensor::{backend::tch::TchTensor, ops::*};
+use crate::TchTensor;
+use burn_tensor::ops::*;
 
 impl<P, const D: usize> Zeros for TchTensor<P, D>
 where

diff --git a/...-tensor/src/tensor/backend/tch/ops/mod.rs → burn-tch/src/ops/mod.rs b/...-tensor/src/tensor/backend/tch/ops/mod.rs → burn-tch/src/ops/mod.rs
diff --git a/burn-tensor/src/tensor/backend/tch/tensor.rs → burn-tch/src/tensor.rs b/burn-tensor/src/tensor/backend/tch/tensor.rs → burn-tch/src/tensor.rs
@@ -1,9 +1,5 @@
-use super::element::TchElement;
-use crate::{
-    backend::{TchBackend, TchDevice},
-    ops::TensorOps,
-    tensor::{Data, Shape},
-};
+use crate::{element::TchElement, TchBackend, TchDevice};
+use burn_tensor::{ops::TensorOps, Data, Shape};
 
 lazy_static::lazy_static! {
     static ref NO_GRAD: tch::NoGradGuard = {
@@ -53,16 +49,6 @@ impl<const D: usize> From<Shape<D>> for TchShape<D> {
     }
 }
 
-impl<const D: usize> From<Vec<i64>> for Shape<D> {
-    fn from(shape: Vec<i64>) -> Self {
-        let mut dims = [1; D];
-        for (i, dim) in shape.into_iter().enumerate() {
-            dims[i] = dim as usize;
-        }
-        Self::new(dims)
-    }
-}
-
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
 pub struct TchKind<P: tch::kind::Element> {
     _p: P,
@@ -99,12 +85,12 @@ impl<P: tch::kind::Element + Default, const D: usize> TchTensor<P, D> {
 #[cfg(test)]
 mod utils {
     use super::*;
-    use crate::{backend::TchBackend, ops::TensorOps};
+    use crate::{backend::TchBackend, element::TchElement};
 
-    impl<P: tch::kind::Element, const D: usize> TchTensor<P, D> {
+    impl<P: TchElement, const D: usize> TchTensor<P, D> {
         pub(crate) fn into_data(self) -> Data<P, D>
         where
-            P: TchElement,
+            P: tch::kind::Element,
         {
             <TchBackend<P> as TensorOps<TchBackend<P>>>::into_data(self)
         }
@@ -131,7 +117,7 @@ impl<P: tch::kind::Element + Default + Copy + std::fmt::Debug, const D: usize> T
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::tensor::Distribution;
+    use burn_tensor::Distribution;
     use rand::prelude::StdRng;
     use rand::SeedableRng;
 

diff --git a/...nsor/src/tensor/backend/tch/tensor_ops.rs → burn-tch/src/tensor_ops.rs b/...nsor/src/tensor/backend/tch/tensor_ops.rs → burn-tch/src/tensor_ops.rs
@@ -1,5 +1,5 @@
 use super::{element::TchElement, TchBackend, TchDevice, TchKind, TchShape, TchTensor};
-use crate::{backend::Backend, ops::TensorOps, Data, ElementConversion, Shape};
+use burn_tensor::{backend::Backend, ops::TensorOps, Data, ElementConversion, Shape};
 use std::ops::{Add, Div, Mul, Range, Sub};
 
 impl<E: TchElement> TensorOps<TchBackend<E>> for TchBackend<E> {

diff --git a/burn-tensor/Cargo.toml b/burn-tensor/Cargo.toml
@@ -20,23 +20,16 @@ all-features = false
 no-default-features = true
 
 [features]
-default = ["tch", "ndarray"]
-tch = ["dep:tch"]
+default = ["ndarray"]
 ndarray = ["dep:ndarray", "dep:libm"]
-doc = ["dep:tch", "tch/doc-only", "dep:ndarray"]
+export_tests = []
 
 [dependencies]
 num-traits = "0.2"
 derive-new = "0.5"
 rand = "0.8"
 half = { version = "1.6", features = ["num-traits"] } # needs to be 1.6 to work with tch
 
-# Backends
-tch = { version = "0.8", optional = true }
-lazy_static = "1.4"
-
-
-
 # NdArray
 ndarray = { version = "0.15", optional = true }
 libm = { version = "0.2", optional = true }

diff --git a/burn-tensor/examples/autodiff_simple.rs b/burn-tensor/examples/autodiff_simple.rs
@@ -43,7 +43,7 @@ fn main() {
     #[cfg(feature = "ndarray")]
     {
         run::<backend::NdArrayBackend<f32>>(x.clone(), y.clone());
-        run_ad::<backend::NdArrayADBackend<f32>>(x.clone(), y.clone());
+        run_ad::<backend::NdArrayADBackend<f32>>(x, y);
     }
 
     #[cfg(feature = "tch")]

diff --git a/burn-tensor/src/lib.rs b/burn-tensor/src/lib.rs
@@ -6,5 +6,8 @@ pub use graph::grad::Gradients;
 
 mod tensor;
 
+#[cfg(feature = "export_tests")]
+mod tests;
+
 pub use half::f16;
 pub use tensor::*;
diff --git a/burn-tensor/src/tensor/backend/autodiff/backend.rs b/burn-tensor/src/tensor/backend/autodiff/backend.rs
@@ -90,6 +90,3 @@ impl<B: Backend> ADBackend for ADBackendDecorator<B> {
 #[cfg(feature = "ndarray")]
 pub type ADBackendNdArray<E> =
     ADBackendDecorator<crate::tensor::backend::ndarray::NdArrayBackend<E>>;
-
-#[cfg(feature = "tch")]
-pub type ADBackendTch<E> = ADBackendDecorator<crate::tensor::backend::tch::TchBackend<E>>;
diff --git a/burn-tensor/src/tensor/backend/autodiff/tensor/base.rs b/burn-tensor/src/tensor/backend/autodiff/tensor/base.rs
@@ -41,14 +41,4 @@ pub mod helper {
         pub type TestADTensor<E, const D: usize> = Tensor<ADBackendNdArray<E>, D>;
     }
     pub use helper_impl::*;
-
-    #[cfg(feature = "tch")]
-    #[cfg(not(feature = "ndarray"))]
-    mod helper_impl {
-        use crate::tensor::backend::autodiff::ADBackendTch;
-        use crate::tensor::Tensor;
-
-        pub type TestADTensor<E, const D: usize> = Tensor<ADBackendTch<E>, D>;
-    }
-    pub use helper_impl::*;
 }
diff --git a/burn-tensor/src/tensor/backend/mod.rs b/burn-tensor/src/tensor/backend/mod.rs
@@ -7,16 +7,6 @@ pub use base::*;
 
 pub(crate) mod autodiff;
 
-#[cfg(feature = "tch")]
-pub(crate) mod tch;
-#[cfg(feature = "tch")]
-pub type TchADBackend<E> = self::autodiff::ADBackendTch<E>;
-#[cfg(feature = "tch")]
-pub type TchBackend<E> = self::tch::TchBackend<E>;
-
-#[cfg(feature = "tch")]
-pub type TchDevice = self::tch::TchDevice;
-
 #[cfg(feature = "ndarray")]
 pub(crate) mod ndarray;
 #[cfg(feature = "ndarray")]
@@ -25,3 +15,5 @@ pub type NdArrayADBackend<E> = self::autodiff::ADBackendNdArray<E>;
 pub type NdArrayBackend<E> = self::ndarray::NdArrayBackend<E>;
 #[cfg(feature = "ndarray")]
 pub type NdArrayDevice = self::ndarray::NdArrayDevice;
+
+pub use autodiff::ADBackendDecorator;
diff --git a/burn-tensor/src/tensor/backend/tch/element.rs b/burn-tensor/src/tensor/backend/tch/element.rs
diff --git a/burn-tensor/src/tensor/element.rs b/burn-tensor/src/tensor/element.rs
@@ -1,4 +1,5 @@
 use crate::{tensor::ops::*, Distribution};
+use half::f16;
 use num_traits::ToPrimitive;
 use rand::prelude::StdRng;
 
@@ -182,3 +183,14 @@ make_element!(
     convert |elem: &dyn ToPrimitive| elem.to_u8().unwrap(),
     random |distribution: Distribution<u8>, rng: &mut StdRng| distribution.sampler(rng).sample()
 );
+make_element!(
+    ty f16 Precision::Half,
+    zero <f16 as num_traits::Zero>::zero(),
+    one <f16 as num_traits::One>::one(),
+    convert |elem: &dyn ToPrimitive| f16::from_f32(elem.to_f32().unwrap()),
+    random |distribution: Distribution<f16>, rng: &mut StdRng| {
+        let distribution: Distribution<f32> = distribution.convert();
+        let sample = distribution.sampler(rng).sample();
+        f16::from_elem(sample)
+    }
+);
diff --git a/burn-tensor/src/tensor/mod.rs b/burn-tensor/src/tensor/mod.rs
@@ -1,4 +1,4 @@
-pub(crate) mod ops;
+pub mod ops;
 pub(crate) mod stats;
 
 mod base;