diff --git a/mu_map/data/preprocessing.py b/mu_map/data/preprocessing.py
deleted file mode 100644
index dc57c6f1de0980238302c460af2d8d9278f8826b..0000000000000000000000000000000000000000
--- a/mu_map/data/preprocessing.py
+++ /dev/null
@@ -1,38 +0,0 @@
-import torch
-
-
-def norm_max(tensor: torch.Tensor):
-    return (tensor - tensor.min()) / (tensor.max() - tensor.min())
-
-
-class MaxNorm:
-    def __call__(self, tensor: torch.Tensor):
-        return norm_max(tensor)
-
-
-def norm_mean(tensor: torch.Tensor):
-    return tensor / tensor.mean()
-
-
-class MeanNorm:
-    def __call__(self, tensor: torch.Tensor):
-        return norm_mean(tensor)
-
-
-def norm_gaussian(tensor: torch.Tensor):
-    return (tensor - tensor.mean()) / tensor.std()
-
-
-class GaussianNorm:
-    def __call__(self, tensor: torch.Tensor):
-        return norm_gaussian(tensor)
-
-
-__all__ = [
-    norm_max.__name__,
-    norm_mean.__name__,
-    norm_gaussian.__name__,
-    MaxNorm.__name__,
-    MeanNorm.__name__,
-    GaussianNorm.__name__,
-]
diff --git a/mu_map/dataset/__init__.py b/mu_map/dataset/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/mu_map/data/datasets.py b/mu_map/dataset/default.py
similarity index 100%
rename from mu_map/data/datasets.py
rename to mu_map/dataset/default.py
diff --git a/mu_map/data/mock.py b/mu_map/dataset/mock.py
similarity index 98%
rename from mu_map/data/mock.py
rename to mu_map/dataset/mock.py
index 2bb2b9ae60b3506232b5a79061fe6ea3c7071fa6..de564b0f9b224cc334febc338a46c01af907dd82 100644
--- a/mu_map/data/mock.py
+++ b/mu_map/dataset/mock.py
@@ -1,4 +1,4 @@
-from mu_map.data.datasets import MuMapDataset
+from mu_map.dataset.default import MuMapDataset
 
 
 class MuMapMockDataset(MuMapDataset):
diff --git a/mu_map/dataset/normalization.py b/mu_map/dataset/normalization.py
new file mode 100644
index 0000000000000000000000000000000000000000..85ec9afc52c2905182c6d9705b2990706eb5a81f
--- /dev/null
+++ b/mu_map/dataset/normalization.py
@@ -0,0 +1,40 @@
+from torch import Tensor
+
+from mu_map.dataset.transform import Transform
+
+
+def norm_max(tensor: Tensor) -> Tensor:
+    return (tensor - tensor.min()) / (tensor.max() - tensor.min())
+
+
+class MaxNormTransform(Transform):
+    def __call__(self, inputs: Tensor, outputs_expected: Tensor) -> Tuple[Tensor, Tensor]:
+        return norm_max(inputs), outputs_expected
+
+
+def norm_mean(tensor: Tensor):
+    return tensor / tensor.mean()
+
+
+class MeanNormTransform(Transform):
+    def __call__(self, inputs: Tensor, outputs_expected: Tensor) -> Tuple[Tensor, Tensor]:
+        return norm_mean(inputs), outputs_expected
+
+
+def norm_gaussian(tensor: Tensor):
+    return (tensor - tensor.mean()) / tensor.std()
+
+
+class GaussianNormTransform(Transform):
+    def __call__(self, inputs: Tensor, outputs_expected: Tensor) -> Tuple[Tensor, Tensor]:
+        return norm_gaussian(inputs), outputs_expected
+
+
+__all__ = [
+    norm_max.__name__,
+    norm_mean.__name__,
+    norm_gaussian.__name__,
+    MaxNormTransform.__name__,
+    MeanNormTransform.__name__,
+    GaussianNormTransform.__name__,
+]
diff --git a/mu_map/data/patch_dataset.py b/mu_map/dataset/patches.py
similarity index 98%
rename from mu_map/data/patch_dataset.py
rename to mu_map/dataset/patches.py
index 934ce823e6996db784ac424f5c46fe13a0fcbd8d..df9f2c971b991c2278a7e4af9a6376c601af6d71 100644
--- a/mu_map/data/patch_dataset.py
+++ b/mu_map/dataset/patches.py
@@ -4,7 +4,7 @@ import random
 import numpy as np
 import torch
 
-from mu_map.data.datasets import MuMapDataset
+from mu_map.dataset.default import MuMapDataset
 
 
 class MuMapPatchDataset(MuMapDataset):