Refactor/burn core (tracel-ai#143)

djdisodo · Jan 2, 2023 · 54c30ab · 54c30ab
1 parent 2f179f1
commit 54c30ab
Show file tree

Hide file tree

Showing 95 changed files with 212 additions and 340 deletions.
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
@@ -59,14 +59,32 @@ jobs:
       crate: burn-ndarray
     secrets: inherit
 
-  publish-burn:
+
+  publish-burn-core:
     uses: burn-rs/burn/.github/workflows/publish-template.yml@main
     needs:
-      - publish-burn-derive
-      - publish-burn-dataset
       - publish-burn-tensor
       - publish-burn-autodiff
+      - publish-burn-dataset
+      - publish-burn-derive
       - publish-burn-ndarray
+    with:
+      crate: burn-core
+    secrets: inherit
+
+  publish-burn-train:
+    uses: burn-rs/burn/.github/workflows/publish-template.yml@main
+    needs:
+      - publish-burn-core
+    with:
+      crate: burn-train
+    secrets: inherit
+
+  publish-burn:
+    uses: burn-rs/burn/.github/workflows/publish-template.yml@main
+    needs:
+      - publish-burn-core
+      - publish-burn-train
     with:
       crate: burn
     secrets: inherit
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -35,7 +35,12 @@ jobs:
     with:
       crate: burn-autodiff
 
-  test-burn:
+  test-burn-core:
     uses: burn-rs/burn/.github/workflows/test-template.yml@main
     with:
-      crate: burn
+      crate: burn-core
+
+  test-burn-train:
+    uses: burn-rs/burn/.github/workflows/test-template.yml@main
+    with:
+      crate: burn-train
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,8 @@
 [workspace]
 members = [
   "burn",
+  "burn-core",
+  "burn-train",
   "burn-derive",
   "burn-tensor",
   "burn-tensor-testgen",

diff --git a/burn-core/Cargo.toml b/burn-core/Cargo.toml
@@ -0,0 +1,40 @@
+[package]
+name = "burn-core"
+version = "0.4.0"
+authors = ["nathanielsimard <nathaniel.simard.42@gmail.com>"]
+description = "BURN: Burn Unstoppable Rusty Neurons"
+repository = "https://github.com/burn-rs/burn-core"
+readme = "README.md"
+keywords = ["deep-learning", "machine-learning", "tensor", "pytorch", "ndarray"]
+categories = ["science"]
+license = "MIT/Apache-2.0"
+edition = "2021"
+
+[features]
+default = []
+experimental-named-tensor = ["burn-tensor/experimental-named-tensor"]
+
+[dependencies]
+burn-tensor = { version = "0.4.0", path = "../burn-tensor" }
+burn-autodiff = { version = "0.4.0", path = "../burn-autodiff" }
+burn-dataset = { version = "0.4.0", path = "../burn-dataset", default-features = false }
+burn-derive = { version = "0.4.0", path = "../burn-derive" }
+
+# Utilities
+derive-new = "0.5.9"
+rand = "0.8.5"
+log = "0.4.17"
+
+# Serialize Deserialize
+serde = { version = "1.0.151", features = ["derive"] }
+serde_json = "1.0.91"
+flate2 = "1.0.25"
+
+# Parameter & Optimization
+nanoid = "0.4.0"
+
+[dev-dependencies]
+burn-dataset = { version = "0.4.0", path = "../burn-dataset", features = [
+    "fake",
+] }
+burn-ndarray = { version = "0.4.0", path = "../burn-ndarray" }
diff --git a/burn-core/LICENSE-APACHE b/burn-core/LICENSE-APACHE
@@ -0,0 +1 @@
+../LICENSE-APACHE
diff --git a/burn-core/LICENSE-MIT b/burn-core/LICENSE-MIT
@@ -0,0 +1 @@
+../LICENSE-MIT
diff --git a/burn-core/README.md b/burn-core/README.md
@@ -0,0 +1,6 @@
+# Burn Core
+
+This crate should be used with [burn](https://github.com/burn-rs/burn).
+
+[![Current Crates.io Version](https://img.shields.io/crates/v/burn-core.svg)](https://crates.io/crates/burn-core)
+[![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](https://github.com/burn-rs/burn-core/blob/master/README.md)
diff --git a/burn/assets → burn-core/assets b/burn/assets → burn-core/assets
diff --git a/burn/src/config.rs → burn-core/src/config.rs b/burn/src/config.rs → burn-core/src/config.rs
diff --git a/burn/src/data/dataloader/base.rs → burn-core/src/data/dataloader/base.rs b/burn/src/data/dataloader/base.rs → burn-core/src/data/dataloader/base.rs
diff --git a/burn/src/data/dataloader/batch.rs → burn-core/src/data/dataloader/batch.rs b/burn/src/data/dataloader/batch.rs → burn-core/src/data/dataloader/batch.rs
diff --git a/burn/src/data/dataloader/batcher.rs → burn-core/src/data/dataloader/batcher.rs b/burn/src/data/dataloader/batcher.rs → burn-core/src/data/dataloader/batcher.rs
diff --git a/burn/src/data/dataloader/builder.rs → burn-core/src/data/dataloader/builder.rs b/burn/src/data/dataloader/builder.rs → burn-core/src/data/dataloader/builder.rs
diff --git a/burn/src/data/dataloader/mod.rs → burn-core/src/data/dataloader/mod.rs b/burn/src/data/dataloader/mod.rs → burn-core/src/data/dataloader/mod.rs
diff --git a/burn/src/data/dataloader/multithread.rs → burn-core/src/data/dataloader/multithread.rs b/burn/src/data/dataloader/multithread.rs → burn-core/src/data/dataloader/multithread.rs
diff --git a/burn/src/data/dataloader/strategy.rs → burn-core/src/data/dataloader/strategy.rs b/burn/src/data/dataloader/strategy.rs → burn-core/src/data/dataloader/strategy.rs
diff --git a/burn/src/data/mod.rs → burn-core/src/data/mod.rs b/burn/src/data/mod.rs → burn-core/src/data/mod.rs
diff --git a/burn-core/src/lib.rs b/burn-core/src/lib.rs
@@ -0,0 +1,14 @@
+#[macro_use]
+extern crate derive_new;
+
+pub mod config;
+pub mod data;
+pub mod module;
+pub mod nn;
+pub mod optim;
+pub mod tensor;
+
+#[cfg(test)]
+pub type TestBackend = burn_ndarray::NdArrayBackend<f32>;
+#[cfg(test)]
+pub type TestADBackend = burn_autodiff::ADBackendDecorator<TestBackend>;
diff --git a/burn/src/module/base.rs → burn-core/src/module/base.rs b/burn/src/module/base.rs → burn-core/src/module/base.rs
@@ -19,10 +19,15 @@ use burn_tensor::Tensor;
 /// necessary to optimize and train the module on any backend.
 ///
 /// ```rust
-/// use burn::nn;
-/// use burn::module::{Param, Module};
-/// use burn::tensor::Tensor;
-/// use burn::tensor::backend::Backend;
+/// // Not necessary when using the burn crate directly.
+/// use burn_core as burn;
+///
+/// use burn::{
+///     nn,
+///     module::{Param, Module},
+///     tensor::Tensor,
+///     tensor::backend::Backend,
+/// };
 ///
 /// #[derive(Module, Debug)]
 /// struct MyModule<B: Backend> {

diff --git a/burn/src/module/mod.rs → burn-core/src/module/mod.rs b/burn/src/module/mod.rs → burn-core/src/module/mod.rs
diff --git a/burn/src/module/param/base.rs → burn-core/src/module/param/base.rs b/burn/src/module/param/base.rs → burn-core/src/module/param/base.rs
diff --git a/burn/src/module/param/id.rs → burn-core/src/module/param/id.rs b/burn/src/module/param/id.rs → burn-core/src/module/param/id.rs
diff --git a/burn/src/module/param/mod.rs → burn-core/src/module/param/mod.rs b/burn/src/module/param/mod.rs → burn-core/src/module/param/mod.rs
diff --git a/burn/src/module/param/module.rs → burn-core/src/module/param/module.rs b/burn/src/module/param/module.rs → burn-core/src/module/param/module.rs
diff --git a/burn/src/module/param/tensor.rs → burn-core/src/module/param/tensor.rs b/burn/src/module/param/tensor.rs → burn-core/src/module/param/tensor.rs
diff --git a/burn/src/module/param/visitor.rs → burn-core/src/module/param/visitor.rs b/burn/src/module/param/visitor.rs → burn-core/src/module/param/visitor.rs
diff --git a/burn/src/module/state.rs → burn-core/src/module/state.rs b/burn/src/module/state.rs → burn-core/src/module/state.rs
diff --git a/burn/src/nn/attention/mask.rs → burn-core/src/nn/attention/mask.rs b/burn/src/nn/attention/mask.rs → burn-core/src/nn/attention/mask.rs
diff --git a/burn/src/nn/attention/mha.rs → burn-core/src/nn/attention/mha.rs b/burn/src/nn/attention/mha.rs → burn-core/src/nn/attention/mha.rs
diff --git a/burn/src/nn/attention/mod.rs → burn-core/src/nn/attention/mod.rs b/burn/src/nn/attention/mod.rs → burn-core/src/nn/attention/mod.rs
diff --git a/burn/src/nn/cache/autoregressive.rs → burn-core/src/nn/cache/autoregressive.rs b/burn/src/nn/cache/autoregressive.rs → burn-core/src/nn/cache/autoregressive.rs
diff --git a/burn/src/nn/cache/base.rs → burn-core/src/nn/cache/base.rs b/burn/src/nn/cache/base.rs → burn-core/src/nn/cache/base.rs
diff --git a/burn/src/nn/cache/mod.rs → burn-core/src/nn/cache/mod.rs b/burn/src/nn/cache/mod.rs → burn-core/src/nn/cache/mod.rs
diff --git a/burn/src/nn/dropout.rs → burn-core/src/nn/dropout.rs b/burn/src/nn/dropout.rs → burn-core/src/nn/dropout.rs
diff --git a/burn/src/nn/embedding.rs → burn-core/src/nn/embedding.rs b/burn/src/nn/embedding.rs → burn-core/src/nn/embedding.rs
diff --git a/burn/src/nn/gelu.rs → burn-core/src/nn/gelu.rs b/burn/src/nn/gelu.rs → burn-core/src/nn/gelu.rs
diff --git a/burn/src/nn/layer_norm.rs → burn-core/src/nn/layer_norm.rs b/burn/src/nn/layer_norm.rs → burn-core/src/nn/layer_norm.rs
diff --git a/burn/src/nn/linear.rs → burn-core/src/nn/linear.rs b/burn/src/nn/linear.rs → burn-core/src/nn/linear.rs
diff --git a/burn/src/nn/loss/cross_entropy.rs → burn-core/src/nn/loss/cross_entropy.rs b/burn/src/nn/loss/cross_entropy.rs → burn-core/src/nn/loss/cross_entropy.rs
diff --git a/burn/src/nn/loss/mod.rs → burn-core/src/nn/loss/mod.rs b/burn/src/nn/loss/mod.rs → burn-core/src/nn/loss/mod.rs
diff --git a/burn/src/nn/mod.rs → burn-core/src/nn/mod.rs b/burn/src/nn/mod.rs → burn-core/src/nn/mod.rs
diff --git a/burn/src/nn/relu.rs → burn-core/src/nn/relu.rs b/burn/src/nn/relu.rs → burn-core/src/nn/relu.rs
diff --git a/burn/src/nn/transformer/encoder.rs → burn-core/src/nn/transformer/encoder.rs b/burn/src/nn/transformer/encoder.rs → burn-core/src/nn/transformer/encoder.rs
diff --git a/burn/src/nn/transformer/mod.rs → burn-core/src/nn/transformer/mod.rs b/burn/src/nn/transformer/mod.rs → burn-core/src/nn/transformer/mod.rs
diff --git a/burn/src/nn/transformer/pwff.rs → burn-core/src/nn/transformer/pwff.rs b/burn/src/nn/transformer/pwff.rs → burn-core/src/nn/transformer/pwff.rs
diff --git a/burn/src/optim/adam.rs → burn-core/src/optim/adam.rs b/burn/src/optim/adam.rs → burn-core/src/optim/adam.rs
diff --git a/burn/src/optim/base.rs → burn-core/src/optim/base.rs b/burn/src/optim/base.rs → burn-core/src/optim/base.rs
diff --git a/burn/src/optim/decay.rs → burn-core/src/optim/decay.rs b/burn/src/optim/decay.rs → burn-core/src/optim/decay.rs
diff --git a/burn/src/optim/grad_accum.rs → burn-core/src/optim/grad_accum.rs b/burn/src/optim/grad_accum.rs → burn-core/src/optim/grad_accum.rs
diff --git a/burn/src/optim/mod.rs → burn-core/src/optim/mod.rs b/burn/src/optim/mod.rs → burn-core/src/optim/mod.rs
@@ -1,14 +1,14 @@
-pub(super) mod visitor;
-
 pub mod decay;
 pub mod momentum;
 
 mod adam;
 mod base;
 mod grad_accum;
 mod sgd;
+mod visitor;
 
 pub use adam::*;
 pub use base::*;
 pub use grad_accum::*;
 pub use sgd::*;
+pub use visitor::*;
diff --git a/burn/src/optim/momentum.rs → burn-core/src/optim/momentum.rs b/burn/src/optim/momentum.rs → burn-core/src/optim/momentum.rs
diff --git a/burn/src/optim/sgd.rs → burn-core/src/optim/sgd.rs b/burn/src/optim/sgd.rs → burn-core/src/optim/sgd.rs
diff --git a/burn/src/optim/visitor.rs → burn-core/src/optim/visitor.rs b/burn/src/optim/visitor.rs → burn-core/src/optim/visitor.rs
@@ -6,34 +6,35 @@ use burn_tensor::{
     Tensor,
 };
 
+/// Data type that contains gradients for a given backend.
 pub type GradientsParams<B> = TensorContainer<<B as ADBackend>::InnerBackend, ParamId>;
 
 #[derive(new)]
-pub struct GradientsRegister<'a, B: ADBackend, O> {
+pub(crate) struct GradientsRegister<'a, B: ADBackend, O> {
     optimizer: &'a O,
     state: &'a mut StateNamed<B::Elem>,
 }
 
 #[derive(new)]
-pub struct GradientsLoader<'a, B: ADBackend, O> {
+pub(crate) struct GradientsLoader<'a, B: ADBackend, O> {
     optimizer: &'a mut O,
     state: &'a StateNamed<B::Elem>,
 }
 
 #[derive(new)]
-pub struct GradientsParamsConverter<'a, B: ADBackend> {
+pub(crate) struct GradientsParamsConverter<'a, B: ADBackend> {
     grads: B::Gradients,
     grads_params: &'a mut TensorContainer<B::InnerBackend, ParamId>,
 }
 
 #[derive(new)]
-pub struct ModuleTensorUpdater<'a, B: ADBackend, O> {
+pub(crate) struct ModuleTensorUpdater<'a, B: ADBackend, O> {
     optimizer: &'a mut O,
     grads: GradientsParams<B>,
 }
 
 #[derive(new)]
-pub struct GradientsParamsChangeDevice<'a, B: ADBackend> {
+pub(crate) struct GradientsParamsChangeDevice<'a, B: ADBackend> {
     device: B::Device,
     grads: &'a mut GradientsParams<B>,
 }
@@ -77,6 +78,7 @@ impl<'a, B: ADBackend> ModuleVisitor<B> for GradientsParamsChangeDevice<'a, B> {
     }
 }
 
+/// Update the device of each tensor gradients.
 pub fn to_device_grads<M: ADModule>(
     grads: &mut GradientsParams<M::ADBackend>,
     device: <M::Backend as Backend>::Device,

diff --git a/burn/src/tensor.rs → burn-core/src/tensor.rs b/burn/src/tensor.rs → burn-core/src/tensor.rs
diff --git a/burn/tests/derive_config.rs → burn-core/tests/derive_config.rs b/burn/tests/derive_config.rs → burn-core/tests/derive_config.rs
@@ -1,4 +1,5 @@
 use burn::config::Config;
+use burn_core as burn;
 
 #[derive(Config, Debug, PartialEq, Eq)]
 pub struct TestEmptyStructConfig {}

diff --git a/burn/tests/derive_module.rs → burn-core/tests/derive_module.rs b/burn/tests/derive_module.rs → burn-core/tests/derive_module.rs
@@ -1,6 +1,7 @@
 use burn::module::{Module, Param};
 use burn::tensor::backend::Backend;
 use burn::tensor::{Distribution, Shape, Tensor};
+use burn_core as burn;
 
 pub type TestBackend = burn_ndarray::NdArrayBackend<f32>;
 

diff --git a/burn-train/Cargo.toml b/burn-train/Cargo.toml
@@ -0,0 +1,29 @@
+[package]
+name = "burn-train"
+version = "0.4.0"
+authors = ["nathanielsimard <nathaniel.simard.42@gmail.com>"]
+description = "Training crate for burn"
+repository = "https://github.com/burn-rs/burn"
+readme = "README.md"
+keywords = ["deep-learning", "machine-learning", "tensor", "pytorch", "ndarray"]
+categories = ["science"]
+license = "MIT/Apache-2.0"
+edition = "2021"
+
+[dependencies]
+burn-core = { version = "0.4.0", path = "../burn-core" }
+
+# Console
+indicatif = "0.17.2"
+log4rs = "1.2.0"
+log = "0.4.17"
+
+# Metrics
+nvml-wrapper = "0.8.0"
+textplots = "0.8.0"
+rgb = "0.8.34"
+terminal_size = "0.2.3"
+
+# Utilities
+derive-new = "0.5.9"
+serde = { version = "1.0.151", features = ["derive"] }
diff --git a/burn-train/LICENSE-APACHE b/burn-train/LICENSE-APACHE
@@ -0,0 +1 @@
+../LICENSE-APACHE
diff --git a/burn-train/LICENSE-MIT b/burn-train/LICENSE-MIT
@@ -0,0 +1 @@
+../LICENSE-MIT
diff --git a/burn-train/README.md b/burn-train/README.md
@@ -0,0 +1,6 @@
+# Burn Train
+
+This crate should be used with [burn](https://github.com/burn-rs/burn).
+
+[![Current Crates.io Version](https://img.shields.io/crates/v/burn-train.svg)](https://crates.io/crates/burn-train)
+[![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](https://github.com/burn-rs/burn-train/blob/master/README.md)
diff --git a/burn/src/train/callback/async_callback.rs → burn-train/src/callback/async_callback.rs b/burn/src/train/callback/async_callback.rs → burn-train/src/callback/async_callback.rs
diff --git a/burn/src/train/callback/base.rs → burn-train/src/callback/base.rs b/burn/src/train/callback/base.rs → burn-train/src/callback/base.rs
@@ -1,4 +1,4 @@
-use crate::data::dataloader::Progress;
+use burn_core::data::dataloader::Progress;
 
 pub trait LearnerCallback<T, V>: Send {
     fn on_train_item(&mut self, _item: LearnerItem<T>) {}

diff --git a/burn/src/train/callback/mod.rs → burn-train/src/callback/mod.rs b/burn/src/train/callback/mod.rs → burn-train/src/callback/mod.rs
diff --git a/.../src/train/checkpoint/async_checkpoint.rs → ...-train/src/checkpoint/async_checkpoint.rs b/.../src/train/checkpoint/async_checkpoint.rs → ...-train/src/checkpoint/async_checkpoint.rs
@@ -1,6 +1,6 @@
 use super::{Checkpointer, CheckpointerError};
-use crate::module::State;
-use burn_tensor::Element;
+use burn_core::module::State;
+use burn_core::tensor::Element;
 use std::sync::{mpsc, Arc};
 
 enum Message<E> {

diff --git a/burn/src/train/checkpoint/base.rs → burn-train/src/checkpoint/base.rs b/burn/src/train/checkpoint/base.rs → burn-train/src/checkpoint/base.rs
@@ -1,4 +1,4 @@
-use crate::module::{State, StateError};
+use burn_core::module::{State, StateError};
 
 #[derive(Debug)]
 pub enum CheckpointerError {

diff --git a/burn/src/train/checkpoint/file.rs → burn-train/src/checkpoint/file.rs b/burn/src/train/checkpoint/file.rs → burn-train/src/checkpoint/file.rs
@@ -1,6 +1,6 @@
 use super::{Checkpointer, CheckpointerError};
-use crate::module::State;
-use burn_tensor::Element;
+use burn_core::module::State;
+use burn_core::tensor::Element;
 
 pub struct FileCheckpointer<P> {
     directory: String,

diff --git a/burn/src/train/checkpoint/mod.rs → burn-train/src/checkpoint/mod.rs b/burn/src/train/checkpoint/mod.rs → burn-train/src/checkpoint/mod.rs
diff --git a/burn/src/train/learner/base.rs → burn-train/src/learner/base.rs b/burn/src/train/learner/base.rs → burn-train/src/learner/base.rs
@@ -1,8 +1,8 @@
-use crate::module::ADModule;
-use crate::optim::Optimizer;
-use crate::tensor::backend::Backend;
-use crate::train::checkpoint::Checkpointer;
-use crate::train::LearnerCallback;
+use crate::checkpoint::Checkpointer;
+use crate::LearnerCallback;
+use burn_core::module::ADModule;
+use burn_core::optim::Optimizer;
+use burn_core::tensor::backend::Backend;
 
 /// Learner struct encapsulating all components necessary to train a Neural Network model.
 ///

diff --git a/burn/src/train/learner/builder.rs → burn-train/src/learner/builder.rs b/burn/src/train/learner/builder.rs → burn-train/src/learner/builder.rs
@@ -1,15 +1,15 @@
 use super::log::update_log_file;
 use super::Learner;
-use crate::module::ADModule;
-use crate::optim::Optimizer;
-use crate::train::checkpoint::{AsyncCheckpointer, Checkpointer, FileCheckpointer};
-use crate::train::logger::FileMetricLogger;
-use crate::train::metric::dashboard::cli::CLIDashboardRenderer;
-use crate::train::metric::dashboard::Dashboard;
-use crate::train::metric::{Adaptor, Metric, Numeric};
-use crate::train::AsyncTrainerCallback;
-use burn_tensor::backend::ADBackend;
-use burn_tensor::Element;
+use crate::checkpoint::{AsyncCheckpointer, Checkpointer, FileCheckpointer};
+use crate::logger::FileMetricLogger;
+use crate::metric::dashboard::cli::CLIDashboardRenderer;
+use crate::metric::dashboard::Dashboard;
+use crate::metric::{Adaptor, Metric, Numeric};
+use crate::AsyncTrainerCallback;
+use burn_core::module::ADModule;
+use burn_core::optim::Optimizer;
+use burn_core::tensor::backend::ADBackend;
+use burn_core::tensor::Element;
 use std::sync::Arc;
 
 /// Struct to configure and create a [learner](Learner).

diff --git a/burn/src/train/learner/classification.rs → burn-train/src/learner/classification.rs b/burn/src/train/learner/classification.rs → burn-train/src/learner/classification.rs
@@ -1,6 +1,6 @@
-use crate::tensor::backend::Backend;
-use crate::train::metric::{AccuracyInput, Adaptor, LossInput};
-use burn_tensor::Tensor;
+use crate::metric::{AccuracyInput, Adaptor, LossInput};
+use burn_core::tensor::backend::Backend;
+use burn_core::tensor::Tensor;
 
 /// Simple classification output adapted for multiple metrics.
 #[derive(new)]

diff --git a/burn/src/train/learner/log.rs → burn-train/src/learner/log.rs b/burn/src/train/learner/log.rs → burn-train/src/learner/log.rs
diff --git a/burn/src/train/learner/mod.rs → burn-train/src/learner/mod.rs b/burn/src/train/learner/mod.rs → burn-train/src/learner/mod.rs
diff --git a/burn/src/train/learner/step/mod.rs → burn-train/src/learner/step/mod.rs b/burn/src/train/learner/step/mod.rs → burn-train/src/learner/step/mod.rs
diff --git a/burn/src/train/learner/step/train.rs → burn-train/src/learner/step/train.rs b/burn/src/train/learner/step/train.rs → burn-train/src/learner/step/train.rs
@@ -1,9 +1,7 @@
-use crate::{
-    data::dataloader::DataLoaderIterator,
-    module::ADModule,
-    train::{TrainOutput, TrainStep},
+use crate::{TrainOutput, TrainStep};
+use burn_core::{
+    data::dataloader::DataLoaderIterator, module::ADModule, tensor::backend::ADBackend,
 };
-use burn_tensor::backend::ADBackend;
 use std::sync::mpsc::{Receiver, Sender};
 use std::thread::spawn;
 

diff --git a/burn/src/train/learner/train_val.rs → burn-train/src/learner/train_val.rs b/burn/src/train/learner/train_val.rs → burn-train/src/learner/train_val.rs
@@ -1,11 +1,12 @@
 use super::Learner;
-use crate::data::dataloader::DataLoader;
-use crate::module::ADModule;
-use crate::optim::visitor::{convert_grads, to_device_grads, GradientsParams};
-use crate::optim::{GradientsAccumulator, Optimizer};
-use crate::train::train::MultiDevicesTrainStep;
-use crate::train::LearnerItem;
-use burn_tensor::backend::ADBackend;
+use crate::train::MultiDevicesTrainStep;
+use crate::LearnerItem;
+use burn_core::data::dataloader::DataLoader;
+use burn_core::module::ADModule;
+use burn_core::optim::{
+    convert_grads, to_device_grads, GradientsAccumulator, GradientsParams, Optimizer,
+};
+use burn_core::tensor::backend::ADBackend;
 use std::sync::Arc;
 
 pub struct TrainOutput<B: ADBackend, TO> {

diff --git a/burn/src/train/mod.rs → burn-train/src/lib.rs b/burn/src/train/mod.rs → burn-train/src/lib.rs
@@ -1,3 +1,6 @@
+#[macro_use]
+extern crate derive_new;
+
 pub mod checkpoint;
 pub mod logger;
 pub mod metric;

diff --git a/burn/src/train/logger/async_logger.rs → burn-train/src/logger/async_logger.rs b/burn/src/train/logger/async_logger.rs → burn-train/src/logger/async_logger.rs
diff --git a/burn/src/train/logger/base.rs → burn-train/src/logger/base.rs b/burn/src/train/logger/base.rs → burn-train/src/logger/base.rs
diff --git a/burn/src/train/logger/file.rs → burn-train/src/logger/file.rs b/burn/src/train/logger/file.rs → burn-train/src/logger/file.rs
diff --git a/burn/src/train/logger/metric.rs → burn-train/src/logger/metric.rs b/burn/src/train/logger/metric.rs → burn-train/src/logger/metric.rs
@@ -1,6 +1,5 @@
-use crate::train::metric::MetricEntry;
-
 use super::{AsyncLogger, FileLogger, Logger};
+use crate::metric::MetricEntry;
 use std::collections::HashMap;
 
 pub trait MetricLogger: Send {

diff --git a/burn/src/train/logger/mod.rs → burn-train/src/logger/mod.rs b/burn/src/train/logger/mod.rs → burn-train/src/logger/mod.rs