mila-iqia · lebrice · Dec 2, 2024 · Dec 2, 2024 · Dec 2, 2024 · Dec 2, 2024
diff --git a/...t/datamodules/datamodules_test/.gitignore → ...project/datasets/datasets_test/.gitignore b/...t/datamodules/datamodules_test/.gitignore → ...project/datasets/datasets_test/.gitignore
diff --git a/...t_batch/cifar10_algorithm_no_op_test.yaml → ...t_batch/cifar10_algorithm_no_op_test.yaml b/...t_batch/cifar10_algorithm_no_op_test.yaml → ...t_batch/cifar10_algorithm_no_op_test.yaml
diff --git a/..._batch/cifar10_algorithm_no_op_train.yaml → ..._batch/cifar10_algorithm_no_op_train.yaml b/..._batch/cifar10_algorithm_no_op_train.yaml → ..._batch/cifar10_algorithm_no_op_train.yaml
diff --git a/...tch/cifar10_algorithm_no_op_validate.yaml → ...tch/cifar10_algorithm_no_op_validate.yaml b/...tch/cifar10_algorithm_no_op_validate.yaml → ...tch/cifar10_algorithm_no_op_validate.yaml
diff --git a/...h/fashion_mnist_algorithm_no_op_test.yaml → ...h/fashion_mnist_algorithm_no_op_test.yaml b/...h/fashion_mnist_algorithm_no_op_test.yaml → ...h/fashion_mnist_algorithm_no_op_test.yaml
diff --git a/.../fashion_mnist_algorithm_no_op_train.yaml → .../fashion_mnist_algorithm_no_op_train.yaml b/.../fashion_mnist_algorithm_no_op_train.yaml → .../fashion_mnist_algorithm_no_op_train.yaml
diff --git a/...shion_mnist_algorithm_no_op_validate.yaml → ...shion_mnist_algorithm_no_op_validate.yaml b/...shion_mnist_algorithm_no_op_validate.yaml → ...shion_mnist_algorithm_no_op_validate.yaml
diff --git a/...batch/glue_cola_algorithm_no_op_test.yaml → ...batch/glue_cola_algorithm_no_op_test.yaml b/...batch/glue_cola_algorithm_no_op_test.yaml → ...batch/glue_cola_algorithm_no_op_test.yaml
diff --git a/...atch/glue_cola_algorithm_no_op_train.yaml → ...atch/glue_cola_algorithm_no_op_train.yaml b/...atch/glue_cola_algorithm_no_op_train.yaml → ...atch/glue_cola_algorithm_no_op_train.yaml
diff --git a/...h/glue_cola_algorithm_no_op_validate.yaml → ...h/glue_cola_algorithm_no_op_validate.yaml b/...h/glue_cola_algorithm_no_op_validate.yaml → ...h/glue_cola_algorithm_no_op_validate.yaml
diff --git a/..._batch/imagenet_algorithm_no_op_test.yaml → ..._batch/imagenet_algorithm_no_op_test.yaml b/..._batch/imagenet_algorithm_no_op_test.yaml → ..._batch/imagenet_algorithm_no_op_test.yaml
diff --git a/...batch/imagenet_algorithm_no_op_train.yaml → ...batch/imagenet_algorithm_no_op_train.yaml b/...batch/imagenet_algorithm_no_op_train.yaml → ...batch/imagenet_algorithm_no_op_train.yaml
diff --git a/...ch/imagenet_algorithm_no_op_validate.yaml → ...ch/imagenet_algorithm_no_op_validate.yaml b/...ch/imagenet_algorithm_no_op_validate.yaml → ...ch/imagenet_algorithm_no_op_validate.yaml
diff --git a/...rst_batch/mnist_algorithm_no_op_test.yaml → ...rst_batch/mnist_algorithm_no_op_test.yaml b/...rst_batch/mnist_algorithm_no_op_test.yaml → ...rst_batch/mnist_algorithm_no_op_test.yaml
diff --git a/...st_batch/mnist_algorithm_no_op_train.yaml → ...st_batch/mnist_algorithm_no_op_train.yaml b/...st_batch/mnist_algorithm_no_op_train.yaml → ...st_batch/mnist_algorithm_no_op_train.yaml
diff --git a/...batch/mnist_algorithm_no_op_validate.yaml → ...batch/mnist_algorithm_no_op_validate.yaml b/...batch/mnist_algorithm_no_op_validate.yaml → ...batch/mnist_algorithm_no_op_validate.yaml
diff --git a/docs/examples/jax_image_classification.md b/docs/examples/jax_image_classification.md
@@ -41,5 +41,5 @@ pass uses Jax to calculate the gradients, and the weights are updated by a PyTor
 ## Running the example
 
 ```console
-$ python project/main.py algorithm=jax_image_classifier network=jax_cnn datamodule=cifar10
+$ python project/main.py algorithm=jax_image_classifier network=jax_cnn dataset=cifar10
 ```
diff --git a/docs/examples/text_classification.md b/docs/examples/text_classification.md
@@ -1,7 +1,7 @@
 ---
 additional_python_references:
   - project.algorithms.text_classifier
-  - project.datamodules.text.text_classification
+  - project.datasets.text.text_classification
 ---
 
 # Text Classification (⚡ + 🤗)
@@ -27,9 +27,9 @@ It accepts a `TextClassificationDataModule` as input, along with a network.
 ### Datamodule config
 
 ??? note "Click to show the Datamodule config"
-    Source: project/configs/datamodule/glue_cola.yaml
+    Source: project/configs/dataset/glue_cola.yaml
 
-    {{ inline('project/configs/datamodule/glue_cola.yaml', 4) }}
+    {{ inline('project/configs/dataset/glue_cola.yaml', 4) }}
 
 ## Running the example
 

diff --git a/docs/profiling_test.py b/docs/profiling_test.py
@@ -11,12 +11,12 @@
     algorithm_network_config,
     command_line_arguments,
     command_line_overrides,
-    datamodule_config,
+    dataset_config,
     experiment_dictconfig,
 )
 from project.experiment import (
     instantiate_algorithm,
-    instantiate_datamodule,
+    instantiate_dataset,
     instantiate_trainer,
     setup_logging,
 )
@@ -80,7 +80,7 @@
         experiment=profiling \
         algorithm=image_classifier \
         algorithm/network=fcnet \
-        datamodule=mnist \
+        dataset=mnist \
         trainer.logger.wandb.name="FcNet/MNIST baseline with training" \
         trainer.logger.wandb.tags=["CPU/GPU comparison","GPU","MNIST"]
         """,
@@ -122,7 +122,7 @@ def test_notebook_commands_dont_cause_errors(experiment_dictconfig: DictConfig):
     setup_logging(log_level=config.log_level)
     lightning.seed_everything(config.seed, workers=True)
     _trainer = instantiate_trainer(config)
-    datamodule = instantiate_datamodule(config.datamodule)
-    _algorithm = instantiate_algorithm(config.algorithm, datamodule=datamodule)
+    dataset = instantiate_dataset(config.dataset)
+    _algorithm = instantiate_algorithm(config.algorithm, dataset=dataset)
 
     # Note: Here we don't actually do anything with the objects.
diff --git a/project/algorithms/callbacks/__init__.py b/project/algorithms/callbacks/__init__.py
@@ -1,9 +1,7 @@
-from .callback import Callback
 from .classification_metrics import ClassificationMetricsCallback
 from .samples_per_second import MeasureSamplesPerSecondCallback
 
 __all__ = [
-    "Callback",
     "ClassificationMetricsCallback",
     "MeasureSamplesPerSecondCallback",
 ]
diff --git a/project/algorithms/callbacks/callback.py b/project/algorithms/callbacks/callback.py