Improve regression example

README.md

@@ -325,17 +325,16 @@ WGPU (WebGPU): Cross-Platform GPU Backend 🌐
 Based on the most popular and well-supported Rust graphics library, [WGPU](https://wgpu.rs), this
 backend automatically targets Vulkan, OpenGL, Metal, Direct X11/12, and WebGPU, by using the WebGPU
 shading language [WGSL](https://www.w3.org/TR/WGSL/https://www.w3.org/TR/WGSL/), or optionally
-[SPIR-V](https://www.khronos.org/spir/) when targeting Vulkan. It can also be compiled to Web Assembly
-to run in the browser while leveraging the GPU, see
+[SPIR-V](https://www.khronos.org/spir/) when targeting Vulkan. It can also be compiled to Web
+Assembly to run in the browser while leveraging the GPU, see
 [this demo](https://antimora.github.io/image-classification/). For more information on the benefits
 of this backend, see [this blog](https://burn.dev/blog/cross-platform-gpu-backend).
 
 The WGPU backend is our first "in-house backend", which means we have complete control over its
 implementation details. It is fully optimized with the
 [performance characteristics mentioned earlier](#performance), as it serves as our research
-playground for a variety of optimizations.
-We've since added CUDA, ROCm and SPIR-V support using the same compiler infrastructure, so a kernel
-written for burn once, can run anywhere.
+playground for a variety of optimizations. We've since added CUDA, ROCm and SPIR-V support using the
+same compiler infrastructure, so a kernel written for burn once, can run anywhere.
 
 See the [WGPU Backend README](./crates/burn-wgpu/README.md) and
 [CUDA Backend README](./crates/burn-cuda/README.md) for more details.
@@ -486,9 +485,9 @@ The Burn Book 🔥
 
 To begin working effectively with Burn, it is crucial to understand its key components and
 philosophy. This is why we highly recommend new users to read the first sections of
-[The Burn Book 🔥](https://burn.dev/burn-book/). It provides detailed examples and explanations covering
-every facet of the framework, including building blocks like tensors, modules, and optimizers, all
-the way to advanced usage, like coding your own GPU kernels.
+[The Burn Book 🔥](https://burn.dev/burn-book/). It provides detailed examples and explanations
+covering every facet of the framework, including building blocks like tensors, modules, and
+optimizers, all the way to advanced usage, like coding your own GPU kernels.
 
 > The project is constantly evolving, and we try as much as possible to keep the book up to date
 > with new additions. However, we might miss some details sometimes, so if you see something weird,
@@ -545,8 +544,8 @@ Additional examples:
 
 - [Custom CSV Dataset](./examples/custom-csv-dataset) : Implements a dataset to parse CSV data for a
   regression task.
-- [Regression](./examples/simple-regression) : Trains a simple MLP on the CSV dataset for the
-  regression task.
+- [Regression](./examples/simple-regression) : Trains a simple MLP on the California Housing dataset
+  to predict the median house value for a district.
 - [Custom Image Dataset](./examples/custom-image-dataset) : Trains a simple CNN on custom image
   dataset following a simple folder structure.
 - [Custom Renderer](./examples/custom-renderer) : Implements a custom renderer to display the

burn-book/src/examples.md

@@ -74,7 +74,7 @@ The following additional examples are currently available if you want to check t
 | Example                                                                                                   | Description                                                                                                                                                                                  |
 | :-------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | [Custom CSV Dataset](https://github.com/tracel-ai/burn/tree/main/examples/custom-csv-dataset)             | Implements a dataset to parse CSV data for a regression task.                                                                                                                                |
-| [Regression](https://github.com/tracel-ai/burn/tree/main/examples/simple-regression)                      | Trains a simple MLP on the CSV dataset for the regression task.                                                                                                                              |
+| [Regression](https://github.com/tracel-ai/burn/tree/main/examples/simple-regression)                      | Trains a simple MLP on the California Housing dataset to predict the median house value for a district.                                                                                      |
 | [Custom Image Dataset](https://github.com/tracel-ai/burn/tree/main/examples/custom-image-dataset)         | Trains a simple CNN on custom image dataset following a simple folder structure.                                                                                                             |
 | [Custom Renderer](https://github.com/tracel-ai/burn/tree/main/examples/custom-renderer)                   | Implements a custom renderer to display the [`Learner`](./building-blocks/learner.md) progress.                                                                                              |
 | [Image Classification Web](https://github.com/tracel-ai/burn/tree/main/examples/image-classification-web) | Image classification web browser demo using Burn, WGPU and WebAssembly.                                                                                                                      |
@@ -83,7 +83,7 @@ The following additional examples are currently available if you want to check t
 | [Named Tensor](https://github.com/tracel-ai/burn/tree/main/examples/named-tensor)                         | Performs operations with the experimental `NamedTensor` feature.                                                                                                                             |
 | [ONNX Import Inference](https://github.com/tracel-ai/burn/tree/main/examples/onnx-inference)              | Imports an ONNX model pre-trained on MNIST to perform inference on a sample image with Burn.                                                                                                 |
 | [PyTorch Import Inference](https://github.com/tracel-ai/burn/tree/main/examples/pytorch-import)           | Imports a PyTorch model pre-trained on MNIST to perform inference on a sample image with Burn.                                                                                               |
-| [Text Classification](https://github.com/tracel-ai/burn/tree/main/examples/text-classification)           | Trains a text classification transformer model on the AG News or DbPedia datasets. The trained model can then be used to classify a text sample.                                              |
+| [Text Classification](https://github.com/tracel-ai/burn/tree/main/examples/text-classification)           | Trains a text classification transformer model on the AG News or DbPedia datasets. The trained model can then be used to classify a text sample.                                             |
 | [Text Generation](https://github.com/tracel-ai/burn/tree/main/examples/text-generation)                   | Trains a text generation transformer model on the DbPedia dataset.                                                                                                                           |
 
 For more information on each example, see their respective `README.md` file. Be sure to check out

examples/simple-regression/Cargo.toml

@@ -22,3 +22,7 @@ burn = {path = "../../crates/burn", features=["train"]}
 # Serialization
 log = {workspace = true}
 serde = {workspace = true, features = ["std", "derive"]}
+
+# Displaying results
+textplots = "0.8.6"
+rgb = "0.8.27"

examples/simple-regression/README.md

@@ -2,9 +2,12 @@
 
 The example shows you how to:
 
-- Define a custom dataset for regression problems. We implement the  [Diabetes Toy Dataset](https://huggingface.co/datasets/Jayabalambika/toy-diabetes) 
-from HuggingFace hub. The dataset is also available as part of toy regression datasets in sklearn[datasets](https://scikit-learn.org/stable/modules/generated/sklearn.datasets.load_diabetes.html).
-- Create a data pipeline from a raw dataset to a batched fast DataLoader with min-max feature scaling.
+- Define a custom dataset for regression problems. We implement the
+  [California Housing Dataset](https://huggingface.co/datasets/gvlassis/california_housing) from
+  HuggingFace hub. The dataset is also available as part of toy regression datasets in
+  sklearn[datasets](https://scikit-learn.org/stable/datasets/real_world.html#california-housing-dataset).
+- Create a data pipeline from a raw dataset to a batched fast DataLoader with min-max feature
+  scaling.
 - Define a Simple NN model for regression using Burn Modules.
 
 > **Note**  

examples/simple-regression/examples/regression.rs

@@ -1,67 +1,63 @@
+use burn::{backend::Autodiff, tensor::backend::Backend};
+use simple_regression::{inference, training};
+
+static ARTIFACT_DIR: &str = "/tmp/burn-example-regression";
+
 #[cfg(any(
     feature = "ndarray",
     feature = "ndarray-blas-netlib",
     feature = "ndarray-blas-openblas",
     feature = "ndarray-blas-accelerate",
 ))]
 mod ndarray {
-    use burn::backend::{
-        ndarray::{NdArray, NdArrayDevice},
-        Autodiff,
-    };
-    use simple_regression::training;
+    use burn::backend::ndarray::{NdArray, NdArrayDevice};
 
     pub fn run() {
         let device = NdArrayDevice::Cpu;
-        training::run::<Autodiff<NdArray>>(device);
+        super::run::<NdArray>(device.clone());
     }
 }
 
 #[cfg(feature = "tch-gpu")]
 mod tch_gpu {
-    use burn::backend::{
-        libtorch::{LibTorch, LibTorchDevice},
-        Autodiff,
-    };
-    use simple_regression::training;
+    use burn::backend::libtorch::{LibTorch, LibTorchDevice};
 
     pub fn run() {
         #[cfg(not(target_os = "macos"))]
         let device = LibTorchDevice::Cuda(0);
         #[cfg(target_os = "macos")]
         let device = LibTorchDevice::Mps;
 
-        training::run::<Autodiff<LibTorch>>(device);
+        super::run::<LibTorch>(device);
     }
 }
 
 #[cfg(feature = "wgpu")]
 mod wgpu {
-    use burn::backend::{
-        wgpu::{Wgpu, WgpuDevice},
-        Autodiff,
-    };
-    use simple_regression::training;
+    use burn::backend::wgpu::{Wgpu, WgpuDevice};
 
     pub fn run() {
         let device = WgpuDevice::default();
-        training::run::<Autodiff<Wgpu>>(device);
+        super::run::<Wgpu>(device);
     }
 }
 
 #[cfg(feature = "tch-cpu")]
 mod tch_cpu {
-    use burn::backend::{
-        libtorch::{LibTorch, LibTorchDevice},
-        Autodiff,
-    };
+    use burn::backend::libtorch::{LibTorch, LibTorchDevice};
     use simple_regression::training;
     pub fn run() {
         let device = LibTorchDevice::Cpu;
-        training::run::<Autodiff<LibTorch>>(device);
+        super::run::<LibTorch>(device);
     }
 }
 
+/// Train a regression model and predict results on a number of samples.
+pub fn run<B: Backend>(device: B::Device) {
+    training::run::<Autodiff<B>>(ARTIFACT_DIR, device.clone());
+    inference::infer::<B>(ARTIFACT_DIR, device)
+}
+
 fn main() {
     #[cfg(any(
         feature = "ndarray",