Repository: IGNF/myria3d (language: python)

Issue:
Dev 3.0.2 update dockerfile



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/cicd.yaml b/.github/workflows/cicd.yaml
index d987b0b1..bce81c48 100644
--- a/.github/workflows/cicd.yaml
+++ b/.github/workflows/cicd.yaml
@@ -26,14 +26,13 @@ jobs:
     - name: Build docker image
       run: docker build -t myria3d .
 
-    - name: Fast unit tests first
+    - name: Run pytest
       run: >
         docker run
-        -v /var/data/cicd/CICD_github_assets/myria3d_B2V0.5/inputs/:/myria3d/tests/data/large/
         --ipc=host
         myria3d
         python -m
-        pytest -rA -v -m "not slow"
+        pytest -rA -v
         --ignore=actions-runner
 
 
@@ -42,19 +41,20 @@ jobs:
     - name: Example inference run via Docker with inference-time subtiles overlap to smooth-out results.
       run: >
         docker run
-        -v /var/data/cicd/CICD_github_assets/myria3d_B2V0.5/inputs/:/inputs/
-        -v /var/data/cicd/CICD_github_assets/myria3d_B2V0.5/outputs/:/outputs/
+        -v /var/data/cicd/CICD_github_assets/myria3d_V3.0.2/inputs/:/inputs/
+        -v /var/data/cicd/CICD_github_assets/myria3d_V3.0.2/outputs/:/outputs/
         --ipc=host
         --shm-size=2gb
         myria3d
-        python /myria3d/myria3d/predict.py
-        experiment=predict
-        predict.ckpt_path=/inputs/proto151_V0.0_epoch_056_Myria3DV2.3.0.ckpt
+        python /app/myria3d/predict.py
+        --config-path /inputs/
+        --config-name proto151_V0.0_epoch_80_Myria3DV3.0.0_predict_config_Myria3DV3.0.0.yaml
+        predict.ckpt_path=/inputs/proto151_V0.0_epoch_80_Myria3DV3.0.0.ckpt
         predict.src_las=/inputs/792000_6272000_subset_buildings.las
         predict.output_dir=/outputs/
         predict.subtile_overlap=25
         datamodule.batch_size=20
-        hydra.run.dir=/myria3d
+        hydra.run.dir=/app
 
     - name: Check code neatness (linter)
       run: docker run myria3d flake8
diff --git a/.github/workflows/gh-pages.yml b/.github/workflows/gh-pages.yml
index ccbf296b..d6e2986a 100644
--- a/.github/workflows/gh-pages.yml
+++ b/.github/workflows/gh-pages.yml
@@ -49,14 +49,14 @@ jobs:
         uses: actions/cache@v2
         with:
           path: ${{ env.CONDA }}/envs
-          key: conda-${{ runner.os }}--${{ runner.arch }}--${{ steps.get-date.outputs.today }}-${{ hashFiles('setup_env/requirements.yml') }}-${{ hashFiles('setup_env/requirements.txt') }}
+          key: conda-${{ runner.os }}--${{ runner.arch }}--${{ steps.get-date.outputs.today }}-${{ hashFiles('environment.yml') }}
         env:
-          # Increase this value to reset cache if etc/example-environment.yml has not changed
+          # Increase this value to reset cache if env has not changed.
           CACHE_NUMBER: 0
         id: cache
 
       - name: Update environment if there was no cached env.
-        run: mamba env update -n myria3d -f setup_env/requirements.yml
+        run: mamba env update -n myria3d -f environment.yml
         if: steps.cache.outputs.cache-hit != 'true'
 
       - name: List installed packages
diff --git a/.gitignore b/.gitignore
index c1615866..ebe63bb3 100755
--- a/.gitignore
+++ b/.gitignore
@@ -4,8 +4,11 @@ inputs/
 outputs/
 checkpoints/
 runs/
+data/
+
 # large files that are not tracked
 tests/data/large
+tests/data/toy_dataset.hdf5
 
 
 # Byte-compiled / optimized / DLL files
diff --git a/Dockerfile b/Dockerfile
index d1cde16b..df83e760 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,57 +1,48 @@
-FROM nvidia/cuda:10.2-devel-ubuntu18.04
-# An nvidia image seems to be necessary for torch-points-kernel. 
-# Also, a "devel" image seems required for the same library
+FROM nvidia/cuda:11.3.1-base-ubuntu20.04
 
 # set the IGN proxy, otherwise apt-get and other applications don't work
 # Should be commented out outside of IGN
 ENV http_proxy 'http://192.168.4.9:3128/'
 ENV https_proxy 'http://192.168.4.9:3128/'
 
-# set the timezone, otherwise it asks for it... and freezes
-ENV TZ=Europe/Paris
-RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
-
-# Needed to use apt-get afterwards due to CUDA changes described here since April 27, 2022:
-# https://forums.developer.nvidia.com/t/notice-cuda-linux-repository-key-rotation/212772
-# Not the recommpended method, but else we need wget installed afterwards.
-# We changed to 10.2-devel-ubuntu18.04 so that might not be needed. 
-RUN apt-get update && apt-get upgrade -y && apt-get install -y wget
-RUN apt-key del 7fa2af80
-RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub
-
-# all the apt-get installs
-RUN apt-get update && apt-get upgrade -y && apt-get install -y \
-        software-properties-common  \
-        wget                        \
-        git                         \
-        libgl1-mesa-glx libegl1-mesa libxrandr2 libxrandr2 libxss1 libxcursor1 libxcomposite1 libasound2 libxi6 libxtst6   # package needed for anaconda
-
-RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh \
-        && /bin/bash ~/miniconda.sh -b -p /opt/conda \
-        && rm ~/miniconda.sh
-
-ENV PATH /opt/conda/bin:$PATH
-
-# Only copy necessary files to set up the environment, 
-# to use docker caching if requirements files were not updated.
-WORKDIR /setup_env
-COPY ./setup_env/ .
-
-# install mamba to setup the env faster
-RUN conda install -y mamba -n base -c conda-forge
-# Build the environment
-RUN mamba env create -f requirements.yml
-
-# Copy the repository content in /myria3d 
-WORKDIR /myria3d
+# Remove any third-party apt sources to avoid issues with expiring keys.
+RUN rm -f /etc/apt/sources.list.d/*.list
+
+# Install some basic utilities
+RUN apt-get update && apt-get install -y \
+        curl \
+        ca-certificates \
+        sudo \
+        git \
+        bzip2 \
+        libx11-6 \
+        && rm -rf /var/lib/apt/lists/*
+
+# Create a working directory
+RUN mkdir /app
+
+# Set up the Conda environment and make python accessible via PATH.
+ENV CONDA_AUTO_UPDATE_CONDA=false
+ENV PATH=/miniconda:/miniconda/bin:$PATH 
+COPY environment.yml /app/environment.yml
+RUN curl -sLo /miniconda.sh https://repo.continuum.io/miniconda/Miniconda3-py39_4.10.3-Linux-x86_64.sh \
+        && chmod +x /miniconda.sh \
+        && /miniconda.sh -b -p /miniconda \
+        && rm /miniconda.sh \
+        && /miniconda/bin/conda env update -n base -f /app/environment.yml \
+        && rm /app/environment.yml \
+        && /miniconda/bin/conda clean -ya
+
+# Need to export this for torch_geometric to find where cuda is.
+# See https://github.com/pyg-team/pytorch_geometric/issues/2040#issuecomment-766610625
+ENV LD_LIBRARY_PATH="/miniconda/lib/:$LD_LIBRARY_PATH"
+
+# Check succes of environment creation.
+RUN python -c "import torch_geometric;"
+
+# Copy the repository content in /app 
+WORKDIR /app
 COPY . .
 
-# Make RUN commands use the new environment:
-SHELL ["conda", "run", "-n", "myria3d", "/bin/bash", "-c"]
-
-# the entrypoint garanty that all command will be runned in the conda environment
-ENTRYPOINT ["conda",  \   
-        "run", \
-        "-n", \
-        "myria3d"]
-
+# Set the default command to bash for image inspection.
+CMD ["bash"]
diff --git a/configs/callbacks/default.yaml b/configs/callbacks/default.yaml
index 406c8fbd..1f221750 100755
--- a/configs/callbacks/default.yaml
+++ b/configs/callbacks/default.yaml
@@ -17,11 +17,6 @@ lr_monitor:
 log_iou_by_class:
   _target_: myria3d.callbacks.logging_callbacks.LogIoUByClass
   classification_dict: ${dataset_description.classification_dict}
-  interpolator:  # only used at test time
-    _target_: myria3d.models.interpolation.Interpolator
-    interpolation_k: ${predict.interpolation_k}
-    classification_dict: ${dataset_description.classification_dict}
-    probas_to_save: ${predict.probas_to_save}  # replace by a list of string of class names to select specific probas to save
 
 model_checkpoint:
   _target_: pytorch_lightning.callbacks.ModelCheckpoint
diff --git a/configs/config.yaml b/configs/config.yaml
index e52219ac..e619d873 100755
--- a/configs/config.yaml
+++ b/configs/config.yaml
@@ -29,7 +29,7 @@ defaults:
   - datamodule: hdf5_datamodule.yaml
   - dataset_description: 20220607_151_dalles_proto.yaml  # describes input features and classes
   - callbacks: default.yaml # set this to null if you don't want to use callbacks
-  - model: randla_net_model.yaml
+  - model: randla_net_model.yaml  # other option is pyg_randla_net_model
   
   - logger: comet # set logger here or use command line (e.g. `python run.py logger=wandb`)
   - task: default.yaml # set logger here or use command line (e.g. `python run.py logger=wandb`)
diff --git a/configs/datamodule/hdf5_datamodule.yaml b/configs/datamodule/hdf5_datamodule.yaml
index 492a5b03..e430d1dc 100755
--- a/configs/datamodule/hdf5_datamodule.yaml
+++ b/configs/datamodule/hdf5_datamodule.yaml
@@ -23,9 +23,9 @@ subtile_shape: "square"  # "square" or "disk"
 subtile_overlap_train: 0
 subtile_overlap_predict: "${predict.subtile_overlap}"
 
-batch_size: 12
-num_workers: 2
-prefetch_factor: 2
+batch_size: 32
+num_workers: 3
+prefetch_factor: 3
 
 defaults:
   - transforms: default.yaml
diff --git a/configs/datamodule/transforms/augmentations/none.yaml b/configs/datamodule/transforms/augmentations/none.yaml
new file mode 100644
index 00000000..e69de29b
diff --git a/configs/datamodule/transforms/default.yaml b/configs/datamodule/transforms/default.yaml
index bd402139..c5f2e980 100644
--- a/configs/datamodule/transforms/default.yaml
+++ b/configs/datamodule/transforms/default.yaml
@@ -1,6 +1,6 @@
 defaults:
   - preparations: default.yaml
-  - augmentations: light.yaml
+  - augmentations: none.yaml
   - normalizations: default.yaml
 
 # turn the config dict into ListConfig that will be fed directly to a Compose object
diff --git a/configs/datamodule/transforms/preparations/no_random_subsampling.yaml b/configs/datamodule/transforms/preparations/no_random_subsampling.yaml
new file mode 100644
index 00000000..1c7627c0
--- /dev/null
+++ b/configs/datamodule/transforms/preparations/no_random_subsampling.yaml
@@ -0,0 +1,71 @@
+# default preparations with grid sampling and random sampling.
+
+train:
+  TargetTransform:
+    _target_: myria3d.pctl.transforms.transforms.TargetTransform
+    _args_:
+      - ${dataset_description.classification_preprocessing_dict}
+      - ${dataset_description.classification_dict}
+
+  GridSampling:
+    _target_: torch_geometric.transforms.GridSampling
+    _args_:
+      - 0.25
+
+  MinimumNumNodes:
+    _target_: myria3d.pctl.transforms.transforms.MinimumNumNodes
+    _args_:
+      - 300
+
+  Center:
+    _target_: torch_geometric.transforms.Center
+
+eval:
+  TargetTransform:
+    _target_: myria3d.pctl.transforms.transforms.TargetTransform
+    _args_:
+      - ${dataset_description.classification_preprocessing_dict}
+      - ${dataset_description.classification_dict}
+
+  CopyFullPreparedTargets:
+    _target_: myria3d.pctl.transforms.transforms.CopyFullPreparedTargets
+
+  CopyFullPos:
+    _target_: myria3d.pctl.transforms.transforms.CopyFullPos
+
+  GridSampling:
+    _target_: torch_geometric.transforms.GridSampling
+    _args_:
+      - 0.25
+
+  MinimumNumNodes:
+    _target_: myria3d.pctl.transforms.transforms.MinimumNumNodes
+    _args_:
+      - 300
+
+  # For interpolation
+  CopySampledPos:
+    _target_: myria3d.pctl.transforms.transforms.CopySampledPos
+
+  Center:
+    _target_: torch_geometric.transforms.Center
+  
+predict:
+  CopyFullPos:
+    _target_: myria3d.pctl.transforms.transforms.CopyFullPos
+
+  GridSampling:
+    _target_: torch_geometric.transforms.GridSampling
+    _args_:
+      - 0.25
+
+  MinimumNumNodes:
+    _target_: myria3d.pctl.transforms.transforms.MinimumNumNodes
+    _args_:
+      - 300
+
+  CopySampledPos:
+    _target_: myria3d.pctl.transforms.transforms.CopySampledPos
+
+  Center:
+    _target_: torch_geometric.transforms.Center
\ No newline at end of file
diff --git a/configs/experiment/RandLaNet-Legacy-Overfit.yaml b/configs/experiment/RandLaNet-Legacy-Overfit.yaml
new file mode 100755
index 00000000..f35f8b0d
--- /dev/null
+++ b/configs/experiment/RandLaNet-Legacy-Overfit.yaml
@@ -0,0 +1,25 @@
+# @package _global_
+
+# Nota : call "python myria3d/pctl/dataset/toy_dataset.py" to create a toy dataset before running this.
+defaults:
+  - override /datamodule/transforms/augmentations: none.yaml
+
+
+logger:
+  comet:
+    experiment_name: "RandLaNetOverfit"
+
+trainer:
+  min_epochs: 100
+  max_epochs: 100
+  overfit_batches: 1
+  num_sanity_val_steps: 0
+
+datamodule:
+  batch_size: 6
+  num_workers: 2
+  # runtime.cwd is where application is run from e.g. where run.py is.
+  data_dir: "${hydra:runtime.cwd}/tests/data/"
+  split_csv_path: "${hydra:runtime.cwd}/tests/data/toy_dataset_src/toy_dataset_split.csv"
+  hdf5_file_path: "${hydra:runtime.cwd}/tests/data/toy_dataset.hdf5"
+
diff --git a/configs/experiment/RandLaNet-PyG-Overfit-NoRS.yaml b/configs/experiment/RandLaNet-PyG-Overfit-NoRS.yaml
new file mode 100755
index 00000000..5d298370
--- /dev/null
+++ b/configs/experiment/RandLaNet-PyG-Overfit-NoRS.yaml
@@ -0,0 +1,25 @@
+# @package _global_
+
+# Nota : call "python myria3d/pctl/dataset/toy_dataset.py" to create a toy dataset before running this.
+defaults:
+  - override /model: pyg_randla_net_model.yaml
+  - override /datamodule/transforms/preparations: no_random_subsampling.yaml
+
+
+logger:
+  comet:
+    experiment_name: "RandLaNetOverfit"
+
+trainer:
+  min_epochs: 100
+  max_epochs: 100
+  overfit_batches: 1
+  num_sanity_val_steps: 0
+
+datamodule:
+  batch_size: 12
+  num_workers: 2
+  # runtime.cwd is where application is run from e.g. where run.py is.
+  data_dir: "${hydra:runtime.cwd}/tests/data/"
+  split_csv_path: "${hydra:runtime.cwd}/tests/data/toy_dataset_src/toy_dataset_split.csv"
+  hdf5_file_path: "${hydra:runtime.cwd}/tests/data/toy_dataset.hdf5"
diff --git a/configs/experiment/RandLaNet-PyG-Overfit.yaml b/configs/experiment/RandLaNet-PyG-Overfit.yaml
new file mode 100755
index 00000000..e017fbac
--- /dev/null
+++ b/configs/experiment/RandLaNet-PyG-Overfit.yaml
@@ -0,0 +1,25 @@
+# @package _global_
+
+# Nota : call "python myria3d/pctl/dataset/toy_dataset.py" to create a toy dataset before running this.
+defaults:
+  - override /model: pyg_randla_net_model.yaml
+  - override /datamodule/transforms/augmentations: none.yaml
+
+
+logger:
+  comet:
+    experiment_name: "RandLaNetOverfit"
+
+trainer:
+  min_epochs: 100
+  max_epochs: 100
+  overfit_batches: 1
+  num_sanity_val_steps: 0
+
+datamodule:
+  batch_size: 12
+  num_workers: 2
+  # runtime.cwd is where application is run from e.g. where run.py is.
+  data_dir: "${hydra:runtime.cwd}/tests/data/"
+  split_csv_path: "${hydra:runtime.cwd}/tests/data/toy_dataset_src/toy_dataset_split.csv"
+  hdf5_file_path: "${hydra:runtime.cwd}/tests/data/toy_dataset.hdf5"
diff --git a/configs/experiment/RandLaNet_base_run_FR_legacy.yaml b/configs/experiment/RandLaNet_base_run_FR_legacy.yaml
index b411f204..0aec48e5 100755
--- a/configs/experiment/RandLaNet_base_run_FR_legacy.yaml
+++ b/configs/experiment/RandLaNet_base_run_FR_legacy.yaml
@@ -11,7 +11,7 @@ datamodule:
   batch_size: 32
 
 trainer:
-  num_sanity_val_steps: 0
+  num_sanity_val_steps: 2
   min_epochs: 100
   max_epochs: 150
   # gpus: [1]
diff --git a/configs/experiment/RandLaNet_base_run_FR_pyg_randla_net-MultiGPU.yaml b/configs/experiment/RandLaNet_base_run_FR_pyg_randla_net-MultiGPU.yaml
new file mode 100755
index 00000000..a84e3fe3
--- /dev/null
+++ b/configs/experiment/RandLaNet_base_run_FR_pyg_randla_net-MultiGPU.yaml
@@ -0,0 +1,18 @@
+# @package _global_
+defaults:
+  - RandLaNet_base_run_FR_pyg_randla_net.yaml
+
+logger:
+  comet:
+    experiment_name: "Pyg RandLaNet - FR Data - 2xGPUs"
+
+trainer:
+  strategy: ddp_find_unused_parameters_false
+  # Replace by gpu to simulate multi-gpus training.
+  accelerator: gpu
+  num_processes: 2
+  gpus: 2
+  
+datamodule:
+  # You may want to divide batch size by number of gpus to keep number of steps by epoch unchanged.
+  batch_size: 16
\ No newline at end of file
diff --git a/configs/experiment/RandLaNet_base_run_FR_pyg_randla_net-SQRT-ICFW.yaml b/configs/experiment/RandLaNet_base_run_FR_pyg_randla_net-SQRT-ICFW.yaml
new file mode 100755
index 00000000..a7ee4ba5
--- /dev/null
+++ b/configs/experiment/RandLaNet_base_run_FR_pyg_randla_net-SQRT-ICFW.yaml
@@ -0,0 +1,29 @@
+# @package _global_
+defaults:
+  - override /model: pyg_randla_net_model.yaml
+  - override /datamodule/transforms/augmentations: light.yaml
+  - override /model/criterion: WeightedCrossEntropyLoss.yaml
+
+logger:
+  comet:
+    experiment_name: "RandLaNet_base_run_FR_pyg_randla_net-SQRT-ICFW"
+
+dataset_description:
+  class_weights: [0.19,0.08,0.08,0.36,1.13,3.11,2.05]
+
+trainer:
+  num_sanity_val_steps: 2
+  min_epochs: 100
+  max_epochs: 150
+  # gpus: [1]
+
+model:
+  lr_scheduler:
+    cooldown: 5
+
+callbacks:
+  early_stopping:
+    patience: 20
+
+predict:
+  interpolation_k: 10
\ No newline at end of file
diff --git a/configs/experiment/RandLaNet_base_run_FR_WD.yaml b/configs/experiment/RandLaNet_base_run_FR_pyg_randla_net.yaml
similarity index 70%
rename from configs/experiment/RandLaNet_base_run_FR_WD.yaml
rename to configs/experiment/RandLaNet_base_run_FR_pyg_randla_net.yaml
index fec08453..fd677bb7 100755
--- a/configs/experiment/RandLaNet_base_run_FR_WD.yaml
+++ b/configs/experiment/RandLaNet_base_run_FR_pyg_randla_net.yaml
@@ -1,15 +1,11 @@
 # @package _global_
 defaults:
+  - override /model: pyg_randla_net_model.yaml
   - override /datamodule/transforms/augmentations: light.yaml
-  - override /model/optimizer: AdamW.yaml
-
 
 logger:
   comet:
-    experiment_name: "RandLaNet - FR Data"
-
-datamodule:
-  batch_size: 32
+    experiment_name: "Pyg RandLaNet - FR Data"
 
 trainer:
   num_sanity_val_steps: 2
diff --git a/configs/experiment/RandLaNet_base_run_FR_pyg_randla_net_NoRS.yaml b/configs/experiment/RandLaNet_base_run_FR_pyg_randla_net_NoRS.yaml
new file mode 100755
index 00000000..e932c90c
--- /dev/null
+++ b/configs/experiment/RandLaNet_base_run_FR_pyg_randla_net_NoRS.yaml
@@ -0,0 +1,20 @@
+# @package _global_
+defaults:
+  - override /model: pyg_randla_net_model.yaml
+  - override /datamodule/transforms/preparations: no_random_subsampling.yaml
+  - override /datamodule/transforms/augmentations: light.yaml
+
+logger:
+  comet:
+    experiment_name: "Pyg RandLaNet - FR Data"
+
+trainer:
+  num_sanity_val_steps: 2
+  min_epochs: 100
+  max_epochs: 150
+  # gpus: [1]
+
+callbacks:
+  early_stopping:
+    patience: 20
+
diff --git a/configs/model/default.yaml b/configs/model/default.yaml
index 3b3bc349..d93cfd74 100644
--- a/configs/model/default.yaml
+++ b/configs/model/default.yaml
@@ -9,6 +9,10 @@ ckpt_path: null  # str, for resuming training and finetuning.
 neural_net_class_name: ???
 neural_net_hparams: ???
 
+# Interpolation params
+interpolation_k: ${predict.interpolation_k}  # interpolation at eval time
+num_workers: 4  # for knn_interpolate
+
 ## Evaluation metric - partial for triple (train/val/test) init
 iou:
   _target_: functools.partial
diff --git a/configs/model/pyg_randla_net_model.yaml b/configs/model/pyg_randla_net_model.yaml
new file mode 100755
index 00000000..46d43953
--- /dev/null
+++ b/configs/model/pyg_randla_net_model.yaml
@@ -0,0 +1,12 @@
+defaults:
+  - default.yaml
+
+lr: 0.003933709606504788   # 200-long LR-test-range betw 10^-4 and 3.0
+
+neural_net_class_name: "PyGRandLANet"
+neural_net_hparams:
+  num_features: ${model.d_in}  # 3 (xyz) + num of features
+  num_classes: ${model.num_classes}
+  num_neighbors: 16
+  decimation: 4  # divide by decimation for each of the 4 local encoder.
+  return_logits: true  # to use with crossEntropyLoss directly
diff --git a/configs/model/randla_net_model.yaml b/configs/model/randla_net_model.yaml
index abab93f5..9ff765e2 100755
--- a/configs/model/randla_net_model.yaml
+++ b/configs/model/randla_net_model.yaml
@@ -9,6 +9,4 @@ neural_net_hparams:
   d_in: ${model.d_in}  # 3 (xyz) + num of features
   num_neighbors: 16
   decimation: 4  # divide by decimation for each of the 4 local encoder.
-  dropout: 0.5  # Use 0.0 of false to deactivate
-  interpolation_k: ${predict.interpolation_k}  # interpolation at eval time
-  num_workers: 4  # for knn_interpolate
\ No newline at end of file
+  dropout: 0.5  # Use 0.0 of false to deactivate
\ No newline at end of file
diff --git a/configs/predict/default.yaml b/configs/predict/default.yaml
index 38299d6d..3f7a431b 100644
--- a/configs/predict/default.yaml
+++ b/configs/predict/default.yaml
@@ -17,4 +17,4 @@ probas_to_save: "all"
 # This comes with a computing cost as the effective predicted area is multiplied.
 subtile_overlap: 0
 # Higher interpolation_k to use more context
-interpolation_k: 25
\ No newline at end of file
+interpolation_k: 10
\ No newline at end of file
diff --git a/docs/source/apidoc/default_config.yml b/docs/source/apidoc/default_config.yml
index e1b6b53d..a4472eb2 100644
--- a/docs/source/apidoc/default_config.yml
+++ b/docs/source/apidoc/default_config.yml
@@ -9,7 +9,6 @@ trainer:
   min_epochs: 1
   max_epochs: 1
   log_every_n_steps: 1
-  check_val_every_n_epoch: 1
   weights_summary: null
   progress_bar_refresh_rate: 1
   auto_lr_find: false
@@ -79,20 +78,13 @@ datamodule:
           _target_: myria3d.pctl.transforms.transforms.CopySampledPos
         Center:
           _target_: torch_geometric.transforms.Center
-    augmentations:
-      x_flip:
-        _target_: torch_geometric.transforms.RandomFlip
-        _args_:
-        - 0
-        p: 0.5
-      y_flip:
-        _target_: torch_geometric.transforms.RandomFlip
-        _args_:
-        - 1
-        p: 0.5
+    augmentations: {}
     normalizations:
       NullifyLowestZ:
         _target_: myria3d.pctl.transforms.transforms.NullifyLowestZ
+      NormalizePos:
+        _target_: myria3d.pctl.transforms.transforms.NormalizePos
+        subtile_width: ${datamodule.subtile_width}
       StandardizeRGBAndIntensity:
         _target_: myria3d.pctl.transforms.transforms.StandardizeRGBAndIntensity
     augmentations_list: '${oc.dict.values: datamodule.transforms.augmentations}'
@@ -108,7 +100,11 @@ datamodule:
     _target_: functools.partial
     _args_:
     - ${get_method:myria3d.pctl.points_pre_transform.lidar_hd.lidar_hd_pre_transform}
-  pre_filter: null
+  pre_filter:
+    _target_: functools.partial
+    _args_:
+    - ${get_method:myria3d.pctl.dataset.utils.pre_filter_below_n_points}
+    min_num_nodes: 50
   tile_width: 1000
   subtile_width: 50
   subtile_shape: square
@@ -116,7 +112,7 @@ datamodule:
   subtile_overlap_predict: ${predict.subtile_overlap}
   batch_size: 2
   num_workers: 1
-  prefetch_factor: 2
+  prefetch_factor: 3
 dataset_description:
   _convert_: all
   classification_preprocessing_dict:
@@ -155,6 +151,14 @@ dataset_description:
     9: water
     17: bridge
     64: lasting_above
+  class_weights:
+  - 0.25
+  - 0.1
+  - 0.1
+  - 0.5
+  - 2.0
+  - 2.0
+  - 2.0
   d_in: 12
   num_classes: 7
 callbacks:
@@ -170,11 +174,6 @@ callbacks:
   log_iou_by_class:
     _target_: myria3d.callbacks.logging_callbacks.LogIoUByClass
     classification_dict: ${dataset_description.classification_dict}
-    interpolator:
-      _target_: myria3d.models.interpolation.Interpolator
-      interpolation_k: ${predict.interpolation_k}
-      classification_dict: ${dataset_description.classification_dict}
-      probas_to_save: ${predict.probas_to_save}
   model_checkpoint:
     _target_: pytorch_lightning.callbacks.ModelCheckpoint
     monitor: val/loss_epoch
@@ -206,6 +205,9 @@ model:
     patience: 5
     cooldown: 0
     verbose: true
+  criterion:
+    _target_: torch.nn.CrossEntropyLoss
+    label_smoothing: 0.0
   _target_: myria3d.models.model.Model
   d_in: ${dataset_description.d_in}
   num_classes: ${dataset_description.num_classes}
@@ -217,17 +219,14 @@ model:
     num_neighbors: 16
     decimation: 4
     dropout: 0.5
-    interpolation_k: ${predict.interpolation_k}
-    num_workers: 4
+  interpolation_k: ${predict.interpolation_k}
+  num_workers: 4
   iou:
     _target_: functools.partial
     _args_:
     - ${get_method:torchmetrics.JaccardIndex}
     - ${model.num_classes}
     absent_score: 1.0
-  criterion:
-    _target_: torch.nn.CrossEntropyLoss
-    label_smoothing: 0.0
   momentum: 0.9
   monitor: val/loss_epoch
   lr: 0.004566395347136576
@@ -249,4 +248,4 @@ predict:
   gpus: 0
   probas_to_save: all
   subtile_overlap: 0
-  interpolation_k: 25
+  interpolation_k: 10
diff --git a/docs/source/apidoc/myria3d.models.modules.rst b/docs/source/apidoc/myria3d.models.modules.rst
index 2942570a..f9c5e11d 100644
--- a/docs/source/apidoc/myria3d.models.modules.rst
+++ b/docs/source/apidoc/myria3d.models.modules.rst
@@ -1,8 +1,19 @@
 myria3d.models.modules
 ========================================
 
-RandLA-Net
--------------
+(Pytorch-Geometric) RandLA-Net
+--------------------------------------------------
+
+Our own implementation of RandLA-Net, to process full point clouds of different sizes, making the best use of the `pytorch-geometric framework <https://github.com/pyg-team/pytorch_geometric>`_.
+Soon to be added to the framework `as an example <https://github.com/pyg-team/pytorch_geometric/pull/5117>`_.
+
+.. autoclass:: myria3d.models.modules.pyg_randla_net.PyGRandLANet
+   :members:
+
+(Legacy) RandLA-Net
+--------------------------------------------------
+
+An open source implementation from `aRI0U <https://github.com/aRI0U/RandLA-Net-pytorch/>`_. Only accepts fixed-sized point clouds.
 
 .. autoclass:: myria3d.models.modules.randla_net.RandLANet
    :members:
\ No newline at end of file
diff --git a/docs/source/background/general_design.md b/docs/source/background/general_design.md
index f1badc52..942a7ca1 100644
--- a/docs/source/background/general_design.md
+++ b/docs/source/background/general_design.md
@@ -2,16 +2,37 @@
 
 Here are a few challenges relative to training 3D segmentation models for Aerial High Density Lidar, and the strategies we adopt in Myria3D in order to face them.
 
+## Model should be fast, performant, and practical
+
+**Situation**:
+- Since the seminal PointCloud architecture [[Qi 2016](https://arxiv.org/abs/1612.00593)] and the later [[PointNet++](https://arxiv.org/abs/1706.02413)], there were many attempts to improve these architecture which operate direcly on point clouds in a hierarchical fashion. 
+- Our main requirements are:
+  1) Speed of inference, in order to operate at a national scale.
+  2) performances in large-scale outdoor Lidar settings, on e.g. [SemanticKITTI](http://semantic-kitti.org/) and [Semantic 3D](http://semantic3d.net/) benchmarks, by opposition to e.g. the [S3DIS](https://ieeexplore.ieee.org/document/7780539/) benchmark.
+  3) performances on specific, urban classes such as buildings.
+  4) Conceptual simplicity - following a simple hierarchical structure was preferred to more complex architectures.
+
+**Strategy**:
+- We focused on the RandLA-Net architecture [Hu 2020](https://arxiv.org/abs/1911.11236). 
+    - It was specifically conceived by authors with requirements (1) and (2) in mind. 
+    - It meets requirement (2), by setting SOTA results on both Semantic3D and SemanticKITTI datasets. In particular, it takes the first place from the KPConv architecture by [Thomas et al. 2019](https://arxiv.org/abs/1904.08889) on Semantic3D, and is in particular superior on class `buildings`, meeting requirement (3). 
+    - The official RandLA-Net implementation is ~4 times faster than KPConv on SemanticKITTI, meeting requirement (1).
+    - It is conceptually simple, basically a PointNet++ like encoder-decoder architecture, with random samplings and explicit encoding of local geometric structures.
+
 
 ## Subsampling is important to improve point cloud structure
 
 **Situation**:
-- Point Cloud data, and aerial Lidar data, represent rich data with a level of detail that might hinder the detection of  objects with generally simple structures such as buildings, ground, and trees. On the other hand, smaller, more variate objects might benefit from denser point clouds.
-- Another point to consider is that some 3D semantic segmentation architectures - including the RandLa-Net architecture we leverage - need fixed size point cloud. This means that either subsampling or padding are required. This kind of implementation reduces flexibility and is suboptimal, but to our knowledge there are no alternative RandLa-Net implementation in pytorch that can accept different-size point clouds within the same batch. 
+- Point Cloud data, and aerial Lidar data, represent rich data with a level of detail that might hinder the detection of objects with generally simple structures such as buildings and ground. On the other hand, smaller, more intricate objects might benefit from denser point clouds.
+
+- Until V3.0.*, we experimented with a RandLA-Net architecture implemented in pytorch by [aRI0U](https://github.com/aRI0U/RandLA-Net-pytorch/), which needed fixed size point cloud. Subsampling was thus required. This kind of implementation reduces flexibility and is suboptimal. There was no alternative RandLa-Net implementation in pytorch that can accept different-size point clouds within the same batch.
 
 **Strategy**:
 - We leverage torch_geometric [GridSampling](https://pytorch-geometric.readthedocs.io/en/latest/modules/transforms.html#torch_geometric.transforms.GridSampling) and [FixedPoints](https://pytorch-geometric.readthedocs.io/en/latest/modules/transforms.html#torch_geometric.transforms.FixedPoints) to (i) simplify local point structures with a 0.25m resolution, and (ii) get a fixed size point cloud that can be fed to the mmodel. Grid Sampling has the effect of reducing point cloud size by around a third, with most reductions expected to occur in vegetation.
 
+**Next Steps**:
+- In V3.0.*, we implement our own version of RandLA-Net. Using the capabilities of pytorch-geometric, it accepts variable size point clouds, and also follows the authors paper more closely. We also contribute it to the pytorch-geometric community (see this [pull request](https://github.com/pyg-team/pytorch_geometric/pull/5117)) and benefit from their feedback. We are in the process of specifying a configuration and transforms to use fuller point clouds, taking into account memory limitations. If this works, this will become the defaut.
+
 ## Speed is of the essence
 **Situation**:
 
@@ -28,4 +49,4 @@ Here are a few challenges relative to training 3D segmentation models for Aerial
 - Evaluation of models must be reliable in order to compare solutions. For semantic segmentation models on point cloud, this means that performance metrics (i.e. mean and by-class Intersection-over-Union) should be computed based on a confusion matrix that is computed from all points in all point clouds in the test dataset.
 
 **Strategy**:
-- During test and validation phases, we **do** interpolate logits back to the each sample (point cloud) before computing performance metrics. Interestingly, this enable to compare different subsampling approaches and interpolation methods in a robust way. The interpolation step is triggered in `eval` mode only, and is also leveraged during inference.
\ No newline at end of file
+- During test and validation phases, we **do** interpolate logits back to the each sample (point cloud) before computing performance metrics. Interestingly, this enable to compare different subsampling approaches and interpolation methods in a robust way. The interpolation step is triggered in `eval` mode only, and is of course also leveraged during inference.
\ No newline at end of file
diff --git a/docs/source/guides/development.md b/docs/source/guides/development.md
index bcaa7c56..9b619659 100644
--- a/docs/source/guides/development.md
+++ b/docs/source/guides/development.md
@@ -17,13 +17,14 @@ python -m pytest -rA -v
 
 ## Continuous Integration (CI)
 
-New features are developped in ad-hoc branches (e.g. `refactor-requirements`), and merged in the `dev` branch. When ready, `dev` can be merged in `main`.
+New features are developped in ad-hoc branches (e.g. `dev-Vx.y.z-Feature-Name`).
 
-CI tests are run for pull request to merge on either `dev` or `main` branches, and on pushes to `dev`, `main`, and `prod`. The CI workflow builds a docker image, runs linting, and tests the code.
+CI tests are run for pull request to merge on either the `main` branches, and on pushes to `main`, and `prod` branches. The CI workflow builds a docker image, runs linting, and tests the code.
 
 ## Continuous Delivery (CD)
 
-When the event is a push and not a merge request, this means that there was either a direct push to `dev`|`main`|`prod` or that a merge request was accepted. In this case, if the CI workflow passes, the docker image is tagged with the branch name, resulting in e.g. a `myria3d:prod` image that is up to date with the branch content. See [../tutorials/use.md] for how to leverage such image to run the app.
+When the event is a push and not a merge request, this means that there was either a direct push to `main`|`prod` or that a merge request was accepted. In this case, if the CI workflow passes (i.e. tests pass and code is PEP8 compliant), the created docker image is tagged with the branch name, resulting in e.g. a `myria3d:prod` image that is up to date with the branch content. 
+See [../tutorials/use.md] for how to leverage such image to run the app.
 
 Additionnaly, pushes on the `main` branch build this library documentation, which is hosted on Github pages.
 
diff --git a/docs/source/guides/train_new_model.md b/docs/source/guides/train_new_model.md
index 02220f2d..121f7209 100644
--- a/docs/source/guides/train_new_model.md
+++ b/docs/source/guides/train_new_model.md
@@ -22,13 +22,13 @@ python run.py experiment=RandLaNetDebug
 
 ## Training
 
-Define your experiment hyperparameters in an experiment file in the `configs/experiment` folder. You may stem from one of the provided experiment file (e.g. `RandLaNet_base_run_FR.yaml`). In particular, you will need to define `dataset_description` to specify your classification task - see config `20220607_151_dalles_proto.yaml` for an example.
+Define your experiment hyperparameters in an experiment file in the `configs/experiment` folder. You may stem from one of the provided experiment file (e.g. `RandLaNet_base_run_FR_pyg_randla_net.yaml`). In particular, you will need to define `dataset_description` to specify your classification task - see config `20220607_151_dalles_proto.yaml` for an example.
 
 
 To run the full training and validation for French Lidar HD, run:
 
 ```bash
-python run.py experiment=RandLaNet_base_run_FR
+python run.py experiment=RandLaNet_base_run_FR_pyg_randla_net
 ```
 
 After training, you model best checkpoints and hydra config will be saved in a `DATE/TIME/` subfolder of the `LOGS_DIR` you specified, with an associated hydra `config.yaml`.
@@ -38,9 +38,16 @@ After training, you model best checkpoints and hydra config will be saved in a `
 Pytorch Lightning support au [automated learning rate finder](https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html#auto-lr-find), by means of an Learning Rate-range test (see section 3.3 in [this paper](https://arxiv.org/pdf/1506.01186.pdf) for reference). 
 You can perfom this automatically before training by setting `trainer.auto_lr_find=true` when calling training on your dataset. The best learning rate will be logged and results saved as an image, so that you do not need to perform this test more than once.
 
+### Multi-GPUs
+
+Multi-GPUs training is supported. Refer to e.g. experiment file `RandLaNet_base_run_FR_pyg_randla_net-MultiGPU.yaml` for pytorch lightning flags to activate it. 
+Multi-GPUs training effectively reduces training time by the number of GPUs used. Batch size might need to be reduced to keep a constant number of steps per epoch in DDP.
+
 ## Testing the model
 
-To evaluate per-class IoU on unseen, annotated data, run:
+Test will be automatically performed after each training, using best checkpointeded model.
+
+To manually evaluate per-class IoU on the test set, run:
 
 ```bash
 python run.py \
@@ -50,8 +57,13 @@ task.task_name="test" \
 model.ckpt_path={/path/to/checkpoint.ckpt} \
 trainer.gpus={0 for none, [i] to use GPU number i} \
 ```
-`config-path` and `config-name` means you are using the saved configuration from your training, which contains the path to the prepared HDF5 dataset. 
-If you are using the default configuration, you do not need those.
+ARguments `config-path` and `config-name` means you are using the saved configuration from your training, which contains the path to the prepared HDF5 dataset. 
+
+If you are using defaut configurations, you can call test using a custom experiment:
+
+```bash
+python run.py experiment=test
+```
 
 ## Inference
 
diff --git a/docs/source/tutorials/make_predictions.md b/docs/source/tutorials/make_predictions.md
index 9d8cd40c..e9bf14c2 100644
--- a/docs/source/tutorials/make_predictions.md
+++ b/docs/source/tutorials/make_predictions.md
@@ -47,22 +47,31 @@ A docker image encapsulating the virtual environment and application sources can
 
 To run inference: 
 - Mount the needed volumes with the `-v` option.
-- Always set `--ipc=host` to allow multithreading in pytorch dataloader (as mentionned in [Pytorch's README](https://github.com/pytorch/pytorch#using-pre-built-images)). 
+- Always set `--ipc=host` to allow multithreading (used in pytorch dataloader, as mentionned in [Pytorch's README](https://github.com/pytorch/pytorch#using-pre-built-images)). 
 - Increase the shared memory with `--shm-size=2gb` (which should be enough for 1km*1km point French "Lidar HD" clouds).
+- Set `--gpus=all` to make gpus visible to the container if available.
+
+See [docker-pytorch README](https://github.com/anibali/docker-pytorch#running-pytorch-scripts) for more details plus an additional option to specify user id at runtime.
 
 ```bash
 # specify your paths here as needed
-docker run -v {local_inputs}:/inputs/ -v {local_output}:/outputs/ --ipc=host --shm-size=2gb myria3d.predict {...options...}
+docker run \
+-v {local_inputs}:/inputs/ \
+-v {local_output}:/outputs/ \
+--ipc=host \
+--gpus=all \
+--shm-size=2gb \
+myria3d.predict {...config paths & options...}
 ```
 
-## Specific options
+## Additional options for prediction
 
 
-### Output format
+### Output dimensions
 
-By default, the predicted classification is stored in a new PRedictedClassification format, which is not a config parameter at the moment. The probability [entropy](https://en.wikipedia.org/wiki/Entropy_(information_theory)) is also stored, as a proxy for prediction uncertainty.
+By default, the predicted classification is stored in a new `PredictedClassification` LAS dimension. The [entropy](https://en.wikipedia.org/wiki/Entropy_(information_theory)) of probabilities is also stored and used as a limited proxy for prediction uncertainty.
 
-What one can control is for which classes to save the probabilities. This is achieved with a `predict.probas_to_save` config parameter, which can be either the `all` keyword (to save probabilities for all classes) or a list of specific classes (e.g. `predict.probas_to_save=[building,vegetation]` - notice the absence of space between class names).
+One can control for which classes to save the probabilities. This is achieved with a `predict.probas_to_save` config parameter, which can be either the `all` keyword (to save probabilities for all classes) or a list of specific classes (e.g. `predict.probas_to_save=[building,vegetation]` - note the absence of space between class names).
 
 ### Receptive field overlap at inference time
 
diff --git a/docs/source/tutorials/prepare_dataset.md b/docs/source/tutorials/prepare_dataset.md
index a416ea1d..177c6fcb 100644
--- a/docs/source/tutorials/prepare_dataset.md
+++ b/docs/source/tutorials/prepare_dataset.md
@@ -1,15 +1,15 @@
 # Preparing data for training
 
-### Peprocessing functions
+## Peprocessing functions
 
 The loading function is dataset dependant, and is `lidar_hd_pre_transform` by default. The function takes points loaded from a LAS file via pdal as input, and returns a `pytorch_geometric.Data` object following the standard naming convention of `pytorch_geometric`, plus a list of features names for later use in transforms.
 
 It is adapted to the French Lidar HD data provided by IGN (see [the official page](https://geoservices.ign.fr/lidarhd) - link in French). Return number and color information (RGBI) are scaled to 0-1 interval, a NDVI and an average color ((R+G+B)/3) dimension are created, and points that may be occluded (as indicated by higher return number) have their color set to 0.
 
-You may want to implement your own logic (e.g. with custom, additional features) in diretcory `points_pre_transform`. It then needs to be referenced similarly to `lidar_hd_pre_transform`.
+You may want to implement your own logic (e.g. with custom, additional features) in directory `points_pre_transform`. It then needs to be referenced similarly to `lidar_hd_pre_transform`.
 
 
-### Using your own data
+## Using your own data
 
 Input point clouds need to be splitted in subtiles that can be digested by segmentation models. We found that a receptive field of 50m*50m was a good balance between context and memory intensity. For faster training, this split can be done once, to avoid loading large file in memory multiple times.
 
diff --git a/docs/source/tutorials/setup_install.md b/docs/source/tutorials/setup_install.md
index 4c39b242..90a84951 100644
--- a/docs/source/tutorials/setup_install.md
+++ b/docs/source/tutorials/setup_install.md
@@ -7,33 +7,25 @@
 We use [anaconda](https://www.anaconda.com/products/individual) to manage virtual environments. 
 This makes installing pytorch-related libraries way easier than using pure pip installs.
 
-If you want to use a gpu to speed up model training and inference, make sure that the cuda toolkit is installed on your machine.
-
-```bash
-sudo apt install nvidia-cuda-toolkit
-```
-
-The  [bash](https://github.com/IGNF/myria3d/setup_env/) directory contains everything you need to setup a compatible pytorch virtual environement.
-You can get it by by cloning the entire project with:
-
-```bash
-git clone https://github.com/IGNF/myria3d
-cd myria3d
-```
-or by manually downloading its content.
+We enable CUDA-acceleration in pytorch as part of the defaut virtual environment recipe (see below).
 
 ### Environment Installation
 
-To install the environment, run:
+To install the environment, follow these instructions:
+
 ```bash
-source setup_env/setup_env.sh
+# Install mamba to create the environment faster
+conda install -y mamba -n base -c conda-forge
+# Build it with mamba
+mamba env create -f environment.yml
+# activate it
+conda activate myria3d
 ```
 
-If you do have CUDA, [check you CUDA version](https://varhowto.com/check-cuda-version/) and be sure `cudatoolkit` version in `setup_env/requirements.yml` matches yours.
+Nota: if you do have CUDA, [check you CUDA version](https://varhowto.com/check-cuda-version/) and be sure `cudatoolkit` version in `setup_env/requirements.yml` matches yours. Adapt if needed (including sources for torch-geometric's dependencies given as wheels).
 
-If you do not have CUDA installed, you need comment out `cudatoolkit` dependency and `FORCE_CUDA: 1` environement variable in `setup_env/requirements.yml`
 
-Finally, activate the created environmnt by running
+Finally, activate the created environment by running
 
 ```bash
 conda activate myria3d
diff --git a/environment.yml b/environment.yml
new file mode 100644
index 00000000..95612790
--- /dev/null
+++ b/environment.yml
@@ -0,0 +1,67 @@
+name: myria3d
+channels:
+  - conda-forge
+  - anaconda
+dependencies:
+  - python==3.9.*
+  - pip
+  # cudatoolkit to specify the cuda driver in the conda env
+  - conda-forge::cudatoolkit=11.3.1 # single equal sign there, not a typo
+  - numba==0.55.1
+  # --------- data formats --------- #
+  - numpy==1.20
+  - h5py
+  # --------- geo --------- #
+  - pygeos
+  - shapely
+  - python-pdal
+  # --------- loggers --------- #
+  # --------- linters --------- #
+  - pre-commit # hooks for applying linters on commit
+  - black # code formatting
+  - isort # import sorting
+  - flake8 # code analysis
+  # --------- tests --------- #
+  - pytest==7.1.*
+  - coverage==6.3.*
+  - pytest-cov==3.0.*
+  # --------- others --------- #
+  - python-dotenv # loading env variables from .env file
+  - rich # beautiful text formatting in terminal
+  - sh # for running bash commands in some tests
+  # - scikit-learn # used in some callbacks
+  # - seaborn # used in some callbacks
+  # - jupyterlab # better jupyter notebooks
+  - pudb # debugger
+  - pip:
+      # --------- Deep Learning --------- #
+      # Extra index may need to be on first line 
+      - --extra-index-url https://download.pytorch.org/whl/cu113
+      - torch==1.11.*
+      - torchvision
+      - pytorch-lightning==1.5.9
+      - torchmetrics==0.7.*  # Else, pytorch-lightning will install the latest
+      - comet_ml==3.31.*
+      # Wheels for torch-geometric optionnal dependencies
+      - https://data.pyg.org/whl/torch-1.11.0%2Bcu113/torch_cluster-1.6.0-cp39-cp39-linux_x86_64.whl
+      - https://data.pyg.org/whl/torch-1.11.0%2Bcu113/torch_scatter-2.0.9-cp39-cp39-linux_x86_64.whl
+      - https://data.pyg.org/whl/torch-1.11.0%2Bcu113/torch_sparse-0.6.14-cp39-cp39-linux_x86_64.whl
+      - git+https://github.com/pyg-team/pytorch_geometric.git@2.1.0
+      # Nota: if libcusparse.so.11. errors occur, run
+      # export LD_LIBRARY_PATH="/home/user/miniconda/envs/lib:$LD_LIBRARY_PATH" 
+      # see https://github.com/pyg-team/pytorch_geometric/issues/2040#issuecomment-766610625
+      # --------- Visualization --------- #
+      - pandas==1.4.*
+      - matplotlib==3.5.*
+      # --------- hydra configs --------- #
+      - hydra-core==1.1.*
+      - hydra-colorlog==1.1.*
+      # # --------- Documentation --------- #
+      - sphinx==4.5.*
+      - sphinx_rtd_theme==1.0.*
+      - myst_parser==0.17.*
+      - sphinx_paramlinks==0.5.*
+      - recommonmark==0.7.*
+      - sphinxnotes-mock==1.0.0b0  # still a beta
+      - docutils==0.17
+      - rstcheck==3.3.*  # RST Linter
diff --git a/myria3d/callbacks/logging_callbacks.py b/myria3d/callbacks/logging_callbacks.py
index 7088807c..5ee7bc96 100644
--- a/myria3d/callbacks/logging_callbacks.py
+++ b/myria3d/callbacks/logging_callbacks.py
@@ -6,7 +6,6 @@
 import torch
 from torchmetrics import JaccardIndex
 from torchmetrics.functional.classification.jaccard import _jaccard_from_confmat
-from myria3d.models.interpolation import Interpolator
 from myria3d.utils import utils
 
 log = utils.get_logger(__name__)
@@ -22,11 +21,10 @@ class LogIoUByClass(Callback):
     A Callback to log JaccardIndex for each class.
     """
 
-    def __init__(self, classification_dict: Dict[int, str], interpolator: Interpolator):
+    def __init__(self, classification_dict: Dict[int, str]):
         self.classification_names = classification_dict.values()
         self.num_classes = len(classification_dict)
         self.metric = SingleClassIoU
-        self.itp = interpolator
 
     def get_all_iou_by_class_object(self):
         """Get a dict with schema {class_name:iou_for_class_name, ...}"""
diff --git a/myria3d/models/model.py b/myria3d/models/model.py
index 2e557380..af12fb5d 100755
--- a/myria3d/models/model.py
+++ b/myria3d/models/model.py
@@ -3,12 +3,17 @@
 from pytorch_lightning import LightningModule
 from torch import nn
 from torch_geometric.data import Batch
-from myria3d.models.modules.randla_net import RandLANet
+from myria3d.models.modules.pyg_randla_net import PyGRandLANet
+from myria3d.models.modules.randla_net import (
+    RandLANet,
+    get_batch_tensor_by_enumeration,
+)
 from myria3d.utils import utils
+from torch_geometric.nn import knn_interpolate
 
 log = utils.get_logger(__name__)
 
-MODEL_ZOO = [RandLANet]
+MODEL_ZOO = [RandLANet, PyGRandLANet]
 
 
 def get_neural_net_class(class_name: str) -> nn.Module:
@@ -59,10 +64,13 @@ def __init__(self, **kwargs):
         # it also allows to access params with 'self.hparams' attribute
         self.save_hyperparameters()
 
-        neural_net_class = get_neural_net_class(self.hparams.neural_net_class_name)
-        self.model = neural_net_class(self.hparams.neural_net_hparams)
+        neural_net_class = get_neural_net_class(
+            self.hparams.neural_net_class_name
+        )
+        self.model = neural_net_class(**self.hparams.neural_net_hparams)
 
         self.softmax = nn.Softmax(dim=1)
+        self.criterion = self.hparams.criterion
 
     def setup(self, stage: Optional[str]) -> None:
         """Setup stage: prepare to compute IoU and loss."""
@@ -71,8 +79,6 @@ def setup(self, stage: Optional[str]) -> None:
             self.val_iou = self.hparams.iou()
         if stage == "test":
             self.test_iou = self.hparams.iou()
-        if stage != "predict":
-            self.criterion = self.hparams.criterion
 
     def forward(self, batch: Batch) -> torch.Tensor:
         """Forward pass of neural network.
@@ -82,26 +88,35 @@ def forward(self, batch: Batch) -> torch.Tensor:
             and y (targets, optionnal) in (B*N,C) format.
 
         Returns:
+            torch.Tensor (B*N,1): targets
             torch.Tensor (B*N,C): logits
 
         """
         logits = self.model(batch)
-        return logits
-
-    # def step(self, batch: Batch):
-    #     """Model step, including loss computation.
-
-    #     Args:
-    #         batch (torch_geometric.data.Batch): Batch of data including x (features), pos (xyz positions),
-    #         and y (targets, optionnal) in (B*N,C) format.
-
-    #     Returns:
-    #         torch.Tensor (1), torch.Tensor (B*N,C), torch.Tensor (B*N,C), torch.Tensor: loss, logits, targets
-
-    #     """
-    # logits = self.forward(batch)
-    # loss = self.criterion(logits, batch.y)
-    #     return loss, logits
+        if self.training or "copies" not in batch:
+            # In training mode and for validation, we directly optimize on subsampled points, for
+            # 1) Speed of training - because interpolation multiplies a step duration by a 5-10 factor!
+            # 2) data augmentation at the supervision level.
+            return batch.y, logits  # B*N, C
+
+        # During evaluation on test data and inference, we interpolate predictions back to original positions
+        # KNN is way faster on CPU than on GPU by a 3 to 4 factor.
+        logits = logits.cpu()
+        batch_y = get_batch_tensor_by_enumeration(batch.idx_in_original_cloud)
+        logits = knn_interpolate(
+            logits.cpu(),
+            batch.copies["pos_sampled_copy"].cpu(),
+            batch.copies["pos_copy"].cpu(),
+            batch_x=batch.batch.cpu(),
+            batch_y=batch_y.cpu(),
+            k=self.hparams.interpolation_k,
+            num_workers=self.hparams.num_workers,
+        )
+        targets = None  # no targets in inference mode.
+        if "transformed_y_copy" in batch.copies:
+            # eval (test/val).
+            targets = batch.copies["transformed_y_copy"].to(logits.device)
+        return targets, logits
 
     def on_fit_start(self) -> None:
         """On fit start: get the experiment for easier access."""
@@ -121,17 +136,22 @@ def training_step(self, batch: Batch, batch_idx: int) -> dict:
         Returns:
             dict: a dict containing the loss, logits, and targets.
         """
-        logits = self.forward(batch)
-        targets = batch.y
+        targets, logits = self.forward(batch)
         self.criterion = self.criterion.to(logits.device)
         loss = self.criterion(logits, targets)
-        self.log("train/loss", loss, on_step=True, on_epoch=True, prog_bar=False)
+        self.log(
+            "train/loss", loss, on_step=True, on_epoch=True, prog_bar=False
+        )
 
         with torch.no_grad():
             preds = torch.argmax(logits.detach(), dim=1)
             self.train_iou(preds, targets)
         self.log(
-            "train/iou", self.train_iou, on_step=True, on_epoch=True, prog_bar=True
+            "train/iou",
+            self.train_iou,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
         )
         return {"loss": loss, "logits": logits, "targets": targets}
 
@@ -149,8 +169,7 @@ def validation_step(self, batch: Batch, batch_idx: int) -> dict:
             dict: a dict containing the loss, logits, and targets.
 
         """
-        logits = self.forward(batch)
-        targets = batch.copies["transformed_y_copy"].to(logits.device)
+        targets, logits = self.forward(batch)
         self.criterion = self.criterion.to(logits.device)
         loss = self.criterion(logits, targets)
         self.log("val/loss", loss, on_step=True, on_epoch=True)
@@ -158,7 +177,9 @@ def validation_step(self, batch: Batch, batch_idx: int) -> dict:
         preds = torch.argmax(logits.detach(), dim=1)
         self.val_iou = self.val_iou.to(preds.device)
         self.val_iou(preds, targets)
-        self.log("val/iou", self.val_iou, on_step=True, on_epoch=True, prog_bar=True)
+        self.log(
+            "val/iou", self.val_iou, on_step=True, on_epoch=True, prog_bar=True
+        )
         return {"loss": loss, "logits": logits, "targets": targets}
 
     def on_validation_epoch_end(self) -> None:
@@ -181,8 +202,7 @@ def test_step(self, batch: Batch, batch_idx: int):
             dict: Dictionnary with full-cloud predicted logits as well as the full-cloud (transformed) targets.
 
         """
-        logits = self.forward(batch)
-        targets = batch.copies["transformed_y_copy"].to(logits.device)
+        targets, logits = self.forward(batch)
         self.criterion = self.criterion.to(logits.device)
         loss = self.criterion(logits, targets)
         self.log("test/loss", loss, on_step=True, on_epoch=True)
@@ -190,7 +210,13 @@ def test_step(self, batch: Batch, batch_idx: int):
         preds = torch.argmax(logits, dim=1)
         self.test_iou = self.test_iou.to(preds.device)
         self.test_iou(preds, targets)
-        self.log("test/iou", self.test_iou, on_step=False, on_epoch=True, prog_bar=True)
+        self.log(
+            "test/iou",
+            self.test_iou,
+            on_step=False,
+            on_epoch=True,
+            prog_bar=True,
+        )
 
         return {"loss": loss, "logits": logits, "targets": targets}
 
@@ -207,7 +233,7 @@ def predict_step(self, batch: Batch) -> dict:
             dict: Dictionnary with predicted logits as well as input batch.
 
         """
-        logits = self.forward(batch)
+        _, logits = self.forward(batch)
         return {"logits": logits.detach().cpu()}
 
     def configure_optimizers(self):
@@ -219,7 +245,8 @@ def configure_optimizers(self):
         """
         self.lr = self.hparams.lr  # aliasing for Lightning auto_find_lr
         optimizer = self.hparams.optimizer(
-            params=filter(lambda p: p.requires_grad, self.parameters()), lr=self.lr
+            params=filter(lambda p: p.requires_grad, self.parameters()),
+            lr=self.lr,
         )
         if self.hparams.lr_scheduler is None:
             return optimizer
diff --git a/myria3d/models/modules/pyg_randla_net.py b/myria3d/models/modules/pyg_randla_net.py
new file mode 100644
index 00000000..8c7477ed
--- /dev/null
+++ b/myria3d/models/modules/pyg_randla_net.py
@@ -0,0 +1,351 @@
+import os.path as osp
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import Sequential, Linear
+from tqdm import tqdm
+import torch_geometric.transforms as T
+from torch_geometric.loader import DataLoader
+from torch_geometric.nn import MLP, global_max_pool
+from torch_geometric.utils import softmax
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.nn.pool import knn_graph
+from torch_geometric.nn.unpool import knn_interpolate
+
+from torch_scatter import scatter
+from torchmetrics.functional import jaccard_index
+from torch_geometric.datasets import ShapeNet
+
+
+class PyGRandLANet(torch.nn.Module):
+    def __init__(
+        self,
+        num_features: int,
+        num_classes: int,
+        decimation: int = 4,
+        num_neighbors: int = 16,
+        return_logits: bool = False,
+    ):
+        super().__init__()
+        # An option to return logits instead of probas
+        self.decimation = decimation
+        self.return_logits = return_logits
+
+        # Authors use 8, which might become a bottleneck if num_classes>8 or num_features>8,
+        # and even for the final MLP.
+        d_bottleneck = max(32, num_classes, num_features)
+
+        self.fc0 = Linear(num_features, d_bottleneck)
+        self.block1 = DilatedResidualBlock(num_neighbors, d_bottleneck, 32)
+        self.block2 = DilatedResidualBlock(num_neighbors, 32, 128)
+        self.block3 = DilatedResidualBlock(num_neighbors, 128, 256)
+        self.block4 = DilatedResidualBlock(num_neighbors, 256, 512)
+        self.mlp_summit = SharedMLP([512, 512])
+        self.fp4 = FPModule(1, SharedMLP([512 + 256, 256]))
+        self.fp3 = FPModule(1, SharedMLP([256 + 128, 128]))
+        self.fp2 = FPModule(1, SharedMLP([128 + 32, 32]))
+        self.fp1 = FPModule(1, SharedMLP([32 + 32, d_bottleneck]))
+        self.mlp_end = Sequential(
+            SharedMLP([d_bottleneck, 64, 32], dropout=[0.0, 0.5]),
+            Linear(32, num_classes),
+        )
+
+    def forward(self, batch):
+        ptr = batch.ptr.clone()
+        pos_and_x = torch.cat([batch.pos, batch.x], axis=1)
+        in_0 = (self.fc0(pos_and_x), batch.pos, batch.batch)
+
+        b1_out = self.block1(*in_0)
+        b1_out_decimated, ptr = decimate(b1_out, ptr, self.decimation)
+
+        b2_out = self.block2(*b1_out_decimated)
+        b2_out_decimated, ptr = decimate(b2_out, ptr, self.decimation)
+
+        b3_out = self.block3(*b2_out_decimated)
+        b3_out_decimated, ptr = decimate(b3_out, ptr, self.decimation)
+
+        b4_out = self.block4(*b3_out_decimated)
+        b4_out_decimated, ptr = decimate(b4_out, ptr, self.decimation)
+
+        mlp_out = (
+            self.mlp_summit(b4_out_decimated[0]),
+            b4_out_decimated[1],
+            b4_out_decimated[2],
+        )
+
+        fp4_out = self.fp4(*mlp_out, *b3_out_decimated)
+        fp3_out = self.fp3(*fp4_out, *b2_out_decimated)
+        fp2_out = self.fp2(*fp3_out, *b1_out_decimated)
+        fp1_out = self.fp1(*fp2_out, *b1_out)
+
+        logits = self.mlp_end(fp1_out[0])
+
+        if self.return_logits:
+            return logits
+
+        probas = logits.log_softmax(dim=-1)
+
+        return probas
+
+
+# Default activation, BatchNorm, and resulting MLP used by RandLA-Net authors
+lrelu02_kwargs = {"negative_slope": 0.2}
+
+
+bn099_kwargs = {"momentum": 0.01, "eps": 1e-6}
+
+
+class SharedMLP(MLP):
+    """SharedMLP with new defauts BN and Activation following tensorflow implementation."""
+
+    # def __init__(self, *args, **kwargs):
+    #     # BN + Act always active even at last layer.
+    #     kwargs["plain_last"] = False
+    #     # LeakyRelu with 0.2 slope by default.
+    #     kwargs["act"] = kwargs.get("act", "LeakyReLU")
+    #     kwargs["act_kwargs"] = kwargs.get("act_kwargs", lrelu02_kwargs)
+    #     # BatchNorm with 1-0.99=0.01 momentum and 1e-6 eps by defaut. (pytorch momentum != tensorflow momentum)
+    #     kwargs["norm_kwargs"] = kwargs.get("norm_kwargs", bn099_kwargs)
+    #     super().__init__(*args, **kwargs)
+
+    def __init__(self, *args, act="LeakyReLU", act_kwargs=lrelu02_kwargs, norm_kwargs=bn099_kwargs, **kwargs):
+        # plain_last=False: BN + Act always active even at last layer.
+        # act="LeakyReLU": LeakyRelu with 0.2 slope by default.
+        # norm_kwargs=bn099_kwargs: BatchNorm with 1-0.99=0.01 momentum and 1e-6 eps by defaut. (pytorch momentum != tensorflow momentum)
+        super().__init__(*args, plain_last=False, act=act, act_kwargs=act_kwargs, norm_kwargs=norm_kwargs, **kwargs)
+
+
+class GlobalPooling(torch.nn.Module):
+    """Global Pooling to adapt RandLA-Net to a classification task."""
+
+    def forward(self, x, pos, batch):
+        x = global_max_pool(x, batch)
+        pos = pos.new_zeros((x.size(0), 3))
+        batch = torch.arange(x.size(0), device=batch.device)
+        return x, pos, batch
+
+
+class LocalFeatureAggregation(MessagePassing):
+    """Positional encoding of points in a neighborhood."""
+
+    def __init__(self, d_out, num_neighbors):
+        super().__init__(aggr="add")
+        self.mlp_encoder = SharedMLP([10, d_out // 2])
+        self.mlp_attention = SharedMLP([d_out, d_out], bias=False, act=None, norm=None)
+        self.mlp_post_attention = SharedMLP([d_out, d_out])
+        self.num_neighbors = num_neighbors
+
+    def forward(self, edge_index, x, pos):
+        out = self.propagate(edge_index, x=x, pos=pos)  # N, d_out
+        out = self.mlp_post_attention(out)  # N, d_out
+        return out
+
+    def message(
+        self, x_j: Tensor, pos_i: Tensor, pos_j: Tensor, index: Tensor
+    ) -> Tensor:
+        """
+        Local Spatial Encoding (locSE) and attentive pooling of features.
+        Args:
+            x_j (Tensor): neighboors features (K,d)
+            pos_i (Tensor): centroid position (repeated) (K,3)
+            pos_j (Tensor): neighboors positions (K,3)
+            index (Tensor): index of centroid positions (e.g. [0,...,0,1,...,1,...,N,...,N])
+        returns:
+            (Tensor): locSE weighted by feature attention scores.
+        """
+        # Encode local neighboorhod structural information
+        pos_diff = pos_j - pos_i
+        distance = torch.sqrt((pos_diff * pos_diff).sum(1, keepdim=True))
+        relative_infos = torch.cat(
+            [pos_i, pos_j, pos_diff, distance], dim=1
+        )  # N * K, d
+        local_spatial_encoding = self.mlp_encoder(relative_infos)  # N * K, d
+        local_features = torch.cat([x_j, local_spatial_encoding], dim=1)  # N * K, 2d
+
+        # Attention will weight the different features of x along the neighborhood dimension.
+        att_features = self.mlp_attention(local_features)  # N * K, d_out
+        att_scores = softmax(att_features, index=index)  # N * K, d_out
+
+        return att_scores * local_features  # N * K, d_out
+
+
+class DilatedResidualBlock(MessagePassing):
+    def __init__(
+        self,
+        num_neighbors,
+        d_in: int,
+        d_out: int,
+    ):
+        super().__init__()
+        self.num_neighbors = num_neighbors
+        self.d_in = d_in
+        self.d_out = d_out
+
+        # MLP on input
+        self.mlp1 = SharedMLP([d_in, d_out // 8])
+        # MLP on input, and the result is summed with the output of mlp2
+        self.shortcut = SharedMLP([d_in, d_out], act=None)
+        # MLP on output
+        self.mlp2 = SharedMLP([d_out // 2, d_out], act=None)
+
+        self.lfa1 = LocalFeatureAggregation(d_out // 4, num_neighbors)
+        self.lfa2 = LocalFeatureAggregation(d_out // 2, num_neighbors)
+
+        self.lrelu = torch.nn.LeakyReLU(**lrelu02_kwargs)
+
+    def forward(self, x, pos, batch):
+        edge_index = knn_graph(pos, self.num_neighbors, batch=batch, loop=True)
+
+        shortcut_of_x = self.shortcut(x)  # N, d_out
+        x = self.mlp1(x)  # N, d_out//8
+        x = self.lfa1(edge_index, x, pos)  # N, d_out//2
+        x = self.lfa2(edge_index, x, pos)  # N, d_out//2
+        x = self.mlp2(x)  # N, d_out
+        x = self.lrelu(x + shortcut_of_x)  # N, d_out
+
+        return x, pos, batch
+
+
+def get_decimation_idx(ptr, decimation):
+    """Subsamples each point cloud by a decimation factor.
+    Decimation happens separately for each cloud to prevent emptying point clouds by accident.
+    """
+    batch_size = ptr.size(0) - 1
+    num_nodes = torch.Tensor([ptr[i + 1] - ptr[i] for i in range(batch_size)]).long()
+    decimated_num_nodes = num_nodes // decimation
+    # Always keep at least one node
+    decimated_num_nodes = torch.max(
+        torch.ones_like(decimated_num_nodes), decimated_num_nodes
+    )
+    idx_decim = torch.cat(
+        [
+            (ptr[i] + torch.randperm(decimated_num_nodes[i], device=ptr.device))
+            for i in range(batch_size)
+        ],
+        dim=0,
+    )
+    # Update the ptr for future decimations
+    ptr_decim = ptr.clone()
+    for i in range(batch_size):
+        ptr_decim[i + 1] = ptr_decim[i] + decimated_num_nodes[i]
+    return idx_decim, ptr_decim
+
+
+def decimate(b_out, ptr, decimation):
+    decimation_idx, decimation_ptr = get_decimation_idx(ptr, decimation)
+    b_out_decim = tuple(t[decimation_idx] for t in b_out)
+    return b_out_decim, decimation_ptr
+
+
+class FPModule(torch.nn.Module):
+    """Upsampling with a skip connection."""
+
+    def __init__(self, k, nn):
+        super().__init__()
+        self.k = k
+        self.nn = nn
+
+    def forward(self, x, pos, batch, x_skip, pos_skip, batch_skip):
+        x = knn_interpolate(x, pos, pos_skip, batch, batch_skip, k=self.k)
+        x = torch.cat([x, x_skip], dim=1)
+        x = self.nn(x)
+        return x, pos_skip, batch_skip
+
+
+def main():
+    category = "Airplane"  # Pass in `None` to train on all categories.
+    category_num_classes = 4  # 4 for Airplane - see ShapeNet for details
+    path = osp.join(
+        osp.dirname(osp.realpath(__file__)), "..", "..", "..", "data", "ShapeNet"
+    )
+    transform = T.Compose(
+        [
+            T.RandomJitter(0.01),
+            T.RandomRotate(15, axis=0),
+            T.RandomRotate(15, axis=1),
+            T.RandomRotate(15, axis=2),
+        ]
+    )
+    pre_transform = T.NormalizeScale()
+    train_dataset = ShapeNet(
+        path,
+        category,
+        split="trainval",
+        transform=transform,
+        pre_transform=pre_transform,
+    )
+    test_dataset = ShapeNet(path, category, split="test", pre_transform=pre_transform)
+    train_loader = DataLoader(train_dataset, batch_size=12, shuffle=True, num_workers=6)
+    test_loader = DataLoader(test_dataset, batch_size=12, shuffle=False, num_workers=6)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = PyGRandLANet(3, category_num_classes).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+
+    def train():
+        model.train()
+
+        total_loss = correct_nodes = total_nodes = 0
+        for i, data in tqdm(enumerate(train_loader)):
+            data = data.to(device)
+            optimizer.zero_grad()
+            out = model(data)
+            loss = F.nll_loss(out, data.y)
+            loss.backward()
+            optimizer.step()
+            total_loss += loss.item()
+            correct_nodes += out.argmax(dim=1).eq(data.y).sum().item()
+            total_nodes += data.num_nodes
+
+            if (i + 1) % 10 == 0:
+                print(
+                    f"[{i+1}/{len(train_loader)}] Loss: {total_loss / 10:.4f} "
+                    f"Train Acc: {correct_nodes / total_nodes:.4f}"
+                )
+                total_loss = correct_nodes = total_nodes = 0
+
+    @torch.no_grad()
+    def test(loader):
+        model.eval()
+
+        ious, categories = [], []
+        y_map = torch.empty(loader.dataset.num_classes, device=device).long()
+        for data in loader:
+            data = data.to(device)
+            outs = model(data)
+
+            sizes = (data.ptr[1:] - data.ptr[:-1]).tolist()
+            for out, y, category in zip(
+                outs.split(sizes), data.y.split(sizes), data.category.tolist()
+            ):
+                category = list(ShapeNet.seg_classes.keys())[category]
+                part = ShapeNet.seg_classes[category]
+                part = torch.tensor(part, device=device)
+
+                y_map[part] = torch.arange(part.size(0), device=device)
+
+                iou = jaccard_index(
+                    out[:, part].argmax(dim=-1),
+                    y_map[y],
+                    num_classes=part.size(0),
+                    absent_score=1.0,
+                )
+                ious.append(iou)
+
+            categories.append(data.category)
+
+        iou = torch.tensor(ious, device=device)
+        category = torch.cat(categories, dim=0)
+
+        mean_iou = scatter(iou, category, reduce="mean")  # Per-category IoU.
+        return float(mean_iou.mean())  # Global IoU.
+
+    for epoch in range(1, 31):
+        train()
+        iou = test(test_loader)
+        print(f"Epoch: {epoch:02d}, Test IoU: {iou:.4f}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/myria3d/models/modules/randla_net.py b/myria3d/models/modules/randla_net.py
index 0ea42fbc..36a9a852 100644
--- a/myria3d/models/modules/randla_net.py
+++ b/myria3d/models/modules/randla_net.py
@@ -4,7 +4,6 @@
 import torch.nn as nn
 
 from torch_geometric.nn.pool import knn
-from torch_geometric.nn.unpool import knn_interpolate
 
 
 class RandLANet(nn.Module):
@@ -14,24 +13,18 @@ class RandLANet(nn.Module):
     RandLA-Net: Efficient Semantic Segmentation of Large-Scale Point Clouds
     at https://arxiv.org/abs/1911.11236
 
-    Our modifications:
-    - fc_start = nn.Linear(d_in, d_in * 2) instead of self.fc_start = nn.Linear(d_in, 8) to avoid
-    information bottleneck in cases where d_in is above 8.
-    - Use pytorch-geometric instead of torch_points_kernels for GPU support, and
-    ease of installation of dependencies.
-
     """
 
-    def __init__(self, hparams_net: dict):
+    def __init__(self, **kwargs: dict):
         super(RandLANet, self).__init__()
-        self.d_in = hparams_net.get("d_in", 6)  # xyz + features
-        self.num_neighbors = hparams_net.get("num_neighbors", 16)
-        self.decimation = hparams_net.get("decimation", 4)
-        self.dropout = hparams_net.get("dropout", 0.0)
-        self.num_classes = hparams_net.get("num_classes", 7)
-        self.interpolation_k = hparams_net.get("interpolation_k", 25)
+        self.d_in = kwargs.get("d_in", 6)  # xyz + features
+        self.num_neighbors = kwargs.get("num_neighbors", 16)
+        self.decimation = kwargs.get("decimation", 4)
+        self.dropout = kwargs.get("dropout", 0.0)
+        self.num_classes = kwargs.get("num_classes", 7)
+        self.interpolation_k = kwargs.get("interpolation_k", 25)
 
-        self.num_workers = hparams_net.get("num_workers", 4)
+        self.num_workers = kwargs.get("num_workers", 4)
 
         self.fc_start = nn.Linear(self.d_in, self.d_in * 2)
         self.bn_start = nn.Sequential(
@@ -134,26 +127,7 @@ def forward(self, batch):
         scores = torch.cat(
             [score_cloud.permute(1, 0) for score_cloud in scores]
         )  # B*N, C
-        if self.training or "copies" not in batch:
-            # In training mode and for validation, we directly optimize on subsampled points, for
-            # 1) Speed of training - because interpolation multiplies a step duration by a 5-10 factor!
-            # 2) data augmentation at the supervision level.
-            return scores  # B*N, C
-
-        # During evaluation on test data and inference, we interpolate predictions back to original positions
-        # KNN is way faster on CPU than on GPU by a 3 to 4 factor.
-        scores = scores.cpu()
-        batch_y = get_batch_tensor_by_enumeration(batch.idx_in_original_cloud)
-        scores = knn_interpolate(
-            scores.cpu(),
-            batch.copies["pos_sampled_copy"].cpu(),
-            batch.copies["pos_copy"].cpu(),
-            batch_x=batch.batch.cpu(),
-            batch_y=batch_y.cpu(),
-            k=self.interpolation_k,
-            num_workers=self.num_workers,
-        )
-        return scores  # N1+N2+...+Nn, C
+        return scores
 
     def set_fc_end(self, d_in, dropout, num_classes):
         """Build the final fully connected layer.
diff --git a/myria3d/pctl/dataloader/dataloader.py b/myria3d/pctl/dataloader/dataloader.py
index 6690d3e4..9a9703db 100644
--- a/myria3d/pctl/dataloader/dataloader.py
+++ b/myria3d/pctl/dataloader/dataloader.py
@@ -6,11 +6,14 @@ class GeometricNoneProofDataloader(DataLoader):
     """Torch geometric's dataloader is a simple torch Dataloader with a different Collater.
 
     This overrides the collater with a NoneProof one that will not fail if some Data is None.
-    This
+
     """
 
     def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, collate_fn=GeometricNoneProofCollater(), **kwargs)
+        # Allow overriding - this is needed for pytorch lightning "overfit_batches" function, which
+        # update the dataloader with a different collater.
+        kwargs["collate_fn"] = kwargs.get("collate_fn", GeometricNoneProofCollater())
+        super().__init__(*args, **kwargs)
 
 
 class GeometricNoneProofCollater(Collater):
diff --git a/myria3d/pctl/datamodule/hdf5.py b/myria3d/pctl/datamodule/hdf5.py
index b9ad2ba1..f65719ec 100644
--- a/myria3d/pctl/datamodule/hdf5.py
+++ b/myria3d/pctl/datamodule/hdf5.py
@@ -11,6 +11,7 @@
 
 from myria3d.pctl.dataset.iterable import InferenceDataset
 from myria3d.pctl.dataset.utils import (
+    LAS_PATHS_BY_SPLIT_DICT_TYPE,
     SHAPE_TYPE,
     find_file_in_dir,
     pre_filter_below_n_points,
@@ -31,7 +32,9 @@ def __init__(
         split_csv_path: str,
         hdf5_file_path: str,
         points_pre_transform: Optional[Callable[[ArrayLike], Data]] = None,
-        pre_filter: Optional[Callable[[Data], bool]] = pre_filter_below_n_points,
+        pre_filter: Optional[
+            Callable[[Data], bool]
+        ] = pre_filter_below_n_points,
         tile_width: Number = 1000,
         subtile_width: Number = 50,
         subtile_shape: SHAPE_TYPE = "square",
@@ -45,6 +48,8 @@ def __init__(
         self.split_csv_path = split_csv_path
         self.data_dir = data_dir
         self.hdf5_file_path = hdf5_file_path
+        self._dataset = None  # will be set by self.dataset property
+        self.las_paths_by_split_dict = None  # Will be set from split_csv
 
         self.points_pre_transform = points_pre_transform
         self.pre_filter = pre_filter
@@ -70,8 +75,12 @@ def __init__(
             "preparations_predict_list", []
         )
 
-        self.augmentation_transform: TRANSFORMS_LIST = t.get("augmentations_list", [])
-        self.normalization_transform: TRANSFORMS_LIST = t.get("normalizations_list", [])
+        self.augmentation_transform: TRANSFORMS_LIST = t.get(
+            "augmentations_list", []
+        )
+        self.normalization_transform: TRANSFORMS_LIST = t.get(
+            "normalizations_list", []
+        )
 
     @property
     def train_transform(self) -> Compose:
@@ -83,7 +92,9 @@ def train_transform(self) -> Compose:
 
     @property
     def eval_transform(self) -> Compose:
-        return Compose(self.preparation_eval_transform + self.normalization_transform)
+        return Compose(
+            self.preparation_eval_transform + self.normalization_transform
+        )
 
     @property
     def predict_transform(self) -> Compose:
@@ -91,19 +102,24 @@ def predict_transform(self) -> Compose:
             self.preparation_predict_transform + self.normalization_transform
         )
 
-    def setup(self, stage: Optional[str] = None):
+    def prepare_data(self, stage: Optional[str] = None):
         """Prepare dataset containing train, val, test data."""
+
         if stage in ["fit", "test"] or stage is None:
-            las_files_by_split = None
+            las_paths_by_split_dict: Optional[
+                LAS_PATHS_BY_SPLIT_DICT_TYPE
+            ] = None
             if self.split_csv_path and self.data_dir:
-                las_files_by_split = {}
+                las_paths_by_split_dict = {}
                 split_df = pd.read_csv(self.split_csv_path)
                 for phase in ["train", "val", "test"]:
-                    basenames = split_df[split_df.split == phase].basename.tolist()
-                    las_files_by_split[phase] = [
+                    basenames = split_df[
+                        split_df.split == phase
+                    ].basename.tolist()
+                    las_paths_by_split_dict[phase] = [
                         find_file_in_dir(self.data_dir, b) for b in basenames
                     ]
-                if not len(las_files_by_split):
+                if not len(las_paths_by_split_dict):
                     raise FileNotFoundError(
                         (
                             f"No basename found while parsing directory {self.data_dir}"
@@ -114,18 +130,47 @@ def setup(self, stage: Optional[str] = None):
                 log.warning(
                     "cfg.data_dir and cfg.split_csv_path are both null. Precomputed HDF5 dataset is used."
                 )
-            self.dataset = HDF5Dataset(
-                self.hdf5_file_path,
-                las_files_by_split=las_files_by_split,
-                points_pre_transform=self.points_pre_transform,
-                tile_width=self.tile_width,
-                subtile_width=self.subtile_width,
-                subtile_overlap_train=self.subtile_overlap_train,
-                subtile_shape=self.subtile_shape,
-                pre_filter=self.pre_filter,
-                train_transform=self.train_transform,
-                eval_transform=self.eval_transform,
-            )
+        # Create the dataset in prepare_data, so that it is done one a single GPU.
+        self.las_paths_by_split_dict = las_paths_by_split_dict
+        self.dataset
+
+    # TODO: not needed ?
+    def setup(self, stage: Optional[str] = None) -> None:
+        """Instantiate the (already prepared) dataset (called on all GPUs)."""
+        self.dataset
+
+    @property
+    def dataset(self) -> HDF5Dataset:
+        """Abstraction to ease HDF5 dataset instantiation.
+
+        Args:
+            las_paths_by_split_dict (LAS_PATHS_BY_SPLIT_DICT_TYPE, optional): Maps split (val/train/test) to file path.
+                If specified, the hdf5 file is created at dataset initialization time.
+                Otherwise,a precomputed HDF5 file is used directly without I/O to the HDF5 file.
+                This is usefule for multi-GPU training, where data creation is performed in prepare_data method, and the dataset
+                is then loaded again in each GPU in setup method.
+                Defaults to None.
+
+        Returns:
+            HDF5Dataset: the dataset with train, val, and test data.
+
+        """
+        if self._dataset:
+            return self._dataset
+
+        self._dataset = HDF5Dataset(
+            self.hdf5_file_path,
+            las_paths_by_split_dict=self.las_paths_by_split_dict,
+            points_pre_transform=self.points_pre_transform,
+            tile_width=self.tile_width,
+            subtile_width=self.subtile_width,
+            subtile_overlap_train=self.subtile_overlap_train,
+            subtile_shape=self.subtile_shape,
+            pre_filter=self.pre_filter,
+            train_transform=self.train_transform,
+            eval_transform=self.eval_transform,
+        )
+        return self._dataset
 
     def train_dataloader(self):
         return GeometricNoneProofDataloader(
diff --git a/myria3d/pctl/dataset/copc.py b/myria3d/pctl/dataset/copc.py
index 75c3f7bc..465cc048 100644
--- a/myria3d/pctl/dataset/copc.py
+++ b/myria3d/pctl/dataset/copc.py
@@ -21,7 +21,14 @@
 
 
 class COPCDataset(Dataset):
-    """Dataset for fast data augmentation of large LAS tiles, for deep learning training/inference."""
+    """Dataset for data augmentation of large LAS tiles, for deep learning training/inference, using COPC format.
+    See https://lidarmag.com/2021/12/27/cloud-native-geospatial-lidar-with-the-cloud-optimized-point-cloud/ for more
+    details.
+
+    Nota: the related DataModule is not implemented at the moment.
+    There is a need to validate speed/performance first. Right now, it is not fast enough to support
+    large batch loading for deep learning applications. LAZ decompression occuring in COPC might be a bottleneck.
+    """
 
     def __init__(
         self,
@@ -34,12 +41,16 @@ def __init__(
         if len(tiles_basenames) == 0:
             raise KeyError("Given list of files is empty")
 
-        processed_basenames = [b.replace(".las", ".copc.laz") for b in tiles_basenames]
+        processed_basenames = [
+            b.replace(".las", ".copc.laz") for b in tiles_basenames
+        ]
         self.copc_paths = [osp.join(copc_dir, b) for b in processed_basenames]
 
         if data_dir:
             # CONVERSION TO COPC IF NEEDED
-            raw_paths = [find_file_in_dir(data_dir, b) for b in tiles_basenames]
+            raw_paths = [
+                find_file_in_dir(data_dir, b) for b in tiles_basenames
+            ]
             try:
                 # IndexError if no file is found in dir.
                 [find_file_in_dir(copc_dir, b) for b in processed_basenames]
@@ -51,7 +62,9 @@ def __init__(
                     desc="Conversion to COPC.LAZ format.",
                 ):
                     write_las_to_copc_laz(
-                        las_path, copc_laz_path, add_original_index=add_original_index
+                        las_path,
+                        copc_laz_path,
+                        add_original_index=add_original_index,
                     )
 
     @abstractmethod
@@ -84,7 +97,9 @@ def __getitem__(self, idx):
             data = self.transform(data)
 
         # filter if empty
-        if data is None or (self.pre_filter is not None and self.pre_filter(data)):
+        if data is None or (
+            self.pre_filter is not None and self.pre_filter(data)
+        ):
             return None
 
         return data
@@ -110,7 +125,10 @@ def __init__(
         resolution: float = 0.0,
     ):
         super().__init__(
-            tiles_basenames, copc_dir, data_dir=datadir, add_original_index=False
+            tiles_basenames,
+            copc_dir,
+            data_dir=datadir,
+            add_original_index=False,
         )
 
         self.tile_width = tile_width
@@ -134,8 +152,7 @@ def load_points(self, idx) -> np.ndarray:
         copc_path = self.copc_paths[idx]
         center = get_random_center_in_tile(self.tile_width, self.subtile_width)
         wkt = make_circle_wkt(center, self.subtile_width)
-        reader_kwargs = {"polygon": wkt, "resolution": self.resolution}
-        points = load_from_copc(copc_path, **reader_kwargs)
+        points = load_from_copc(copc_path, polygon=wkt, resolution=self.resolution)
         return points
 
 
@@ -173,7 +190,9 @@ def __init__(
 
         # samples is a list of path-center pairs
         xy_centers = get_mosaic_of_centers(
-            self.tile_width, self.subtile_width, subtile_overlap=subtile_overlap
+            self.tile_width,
+            self.subtile_width,
+            subtile_overlap=subtile_overlap,
         )
         self.samples = []
         for path in self.copc_paths:
@@ -187,8 +206,7 @@ def __len__(self):
     def load_points(self, idx) -> np.ndarray:
         copc_path, center = self.samples[idx]
         wkt = make_circle_wkt(center, self.subtile_width)
-        reader_kwargs = {"polygon": wkt}
-        points = load_from_copc(copc_path, **reader_kwargs)
+        points = load_from_copc(copc_path, polygon=wkt)
         return points
 
 
diff --git a/myria3d/pctl/dataset/hdf5.py b/myria3d/pctl/dataset/hdf5.py
index 2759a943..41039caf 100644
--- a/myria3d/pctl/dataset/hdf5.py
+++ b/myria3d/pctl/dataset/hdf5.py
@@ -14,7 +14,7 @@
 from myria3d.pctl.dataset.utils import (
     SHAPE_TYPE,
     SPLIT_TYPE,
-    LAS_FILES_BY_SPLIT_TYPE,
+    LAS_PATHS_BY_SPLIT_DICT_TYPE,
     pre_filter_below_n_points,
     split_cloud_into_samples,
 )
@@ -29,7 +29,7 @@ class HDF5Dataset(Dataset):
     def __init__(
         self,
         hdf5_file_path: str,
-        las_files_by_split: Optional[LAS_FILES_BY_SPLIT_TYPE] = None,
+        las_paths_by_split_dict: Optional[LAS_PATHS_BY_SPLIT_DICT_TYPE] = None,
         points_pre_transform: Callable = lidar_hd_pre_transform,
         tile_width: Number = 1000,
         subtile_width: Number = 50,
@@ -42,8 +42,8 @@ def __init__(
         """Initialization, taking care of HDF5 dataset preparation if needed, and indexation of its content.
 
         Args:
-            las_files_by_split (Optional[LAS_FILES_BY_SPLIT_TYPE]): should look like
-                las_files_by_split = {'train': ['las1.las','las2.las'], 'val': [...], , 'test': [...]}
+            las_paths_by_split_dict (Optional[LAS_PATHS_BY_SPLIT_DICT_TYPE]): should look like
+                las_paths_by_split_dict = {'train': ['dir/las1.las','dir/las2.las'], 'val': [...], , 'test': [...]}
             hdf5_file_path (str): path to HDF5 dataset
             points_pre_transform (Callable): Function to turn pdal points into a pyg Data object.
             tile_width (Number, optional): width of a LAS tile. Defaults to 1000.
@@ -61,130 +61,89 @@ def __init__(
         self.train_transform = train_transform
         self.eval_transform = eval_transform
 
+        self.tile_width = tile_width
+        self.subtile_width = subtile_width
+        self.subtile_overlap_train = subtile_overlap_train
+        self.subtile_shape = subtile_shape
+
         self.hdf5_file_path = hdf5_file_path
 
-        if las_files_by_split is None:
+        # Instantiates these to null;
+        # They are loaded within __getitem__ to support multi-processing training.
+        self.dataset = None
+        self._samples_hdf5_paths = None
+
+        if las_paths_by_split_dict is None:
             log.warning(
-                "las_files_by_split is None and pre-computed HDF5 dataset is therefore used."
+                "las_paths_by_split_dict is None and pre-computed HDF5 dataset is therefore used."
             )
-            self.index_all_samples()
             return
 
-        # CONVERSION TO a SINGLE H5 file
+        # Add data for all LAS Files into a single hdf5 file.
         os.makedirs(osp.dirname(hdf5_file_path), exist_ok=True)
-        for s, raw_paths in las_files_by_split.items():
+        for split, las_paths in las_paths_by_split_dict.items():
             with h5py.File(hdf5_file_path, "a") as f:
-                if s not in f:
-                    f.create_group(s)
-            for las_path in tqdm(raw_paths, desc=f"Preparing {s} set..."):
-                b = osp.basename(las_path)
-                # Delete incomplete file entry to start from scratch.
-                with h5py.File(hdf5_file_path, "a") as f:
-                    if b in f[s] and "is_complete" not in f[s][b].attrs:
-                        del f[b]
-                # Parse and add subtiles to split group.
-                with h5py.File(hdf5_file_path, "a") as f:
-                    if b not in f[s]:
-                        sample_number = 0
-                        for sample_idx, sample_points in split_cloud_into_samples(
-                            las_path,
-                            tile_width,
-                            subtile_width,
-                            subtile_shape,
-                            subtile_overlap=subtile_overlap_train
-                            if s == "train"
-                            else 0,
-                        ):
-                            data = self.points_pre_transform(sample_points)
-                            if self.pre_filter is not None and self.pre_filter(data):
-                                # e.g. not enough points in this receptive field.
-                                continue
-                            sample_id = str(sample_number).zfill(5)
-                            hdf5_path = osp.join(s, b, sample_id)
-                            hd5f_path_x = osp.join(hdf5_path, "x")
-                            f.create_dataset(
-                                hd5f_path_x,
-                                data.x.shape,
-                                dtype="f",
-                                data=data.x,
-                            )
-                            f[hd5f_path_x].attrs["x_features_names"] = copy.deepcopy(
-                                data.x_features_names
-                            )
-                            f.create_dataset(
-                                osp.join(hdf5_path, "pos"),
-                                data.pos.shape,
-                                dtype="f",
-                                data=data.pos,
-                            )
-                            f.create_dataset(
-                                osp.join(hdf5_path, "y"),
-                                data.y.shape,
-                                dtype="i",
-                                data=data.y,
-                            )
-                            f.create_dataset(
-                                osp.join(hdf5_path, "idx_in_original_cloud"),
-                                sample_idx.shape,
-                                dtype="i",
-                                data=sample_idx,
-                            )
-                            sample_number += 1
-
-                        # A termination flag to report that all samples for this point cloud were included in the df5 file.
-                        f[s][b].attrs["is_complete"] = True
-        self.index_all_samples()
-
-    def index_all_samples(self):
-        self.samples_hdf5_paths = []
-        with h5py.File(self.hdf5_file_path, "a") as f:
-            for s in f.keys():
-                basenames = f[s].keys()
-                for bn in basenames:
-                    self.samples_hdf5_paths += [
-                        osp.join(s, bn, sample_number)
-                        for sample_number in f[s][bn].keys()
-                    ]
-
-    def __len__(self):
-        return len(self.samples_hdf5_paths)
+                if split not in f:
+                    f.create_group(split)
+            for las_path in tqdm(las_paths, desc=f"Preparing {split} set..."):
+                self._add_to_dataset(split, las_path)
 
-    def _get_data(self, sample_hdf5_path) -> Data:
-        with h5py.File(self.hdf5_file_path, "r") as f:
-            grp = f[sample_hdf5_path]
-            # [...] needed to make a copy of content and avoid closing HDF5.
-            # Nota: idx_in_original_cloud SHOULD be np.ndarray, in order to be batched into a list,
-            # which serves to keep track of indivual sample sizes in a simpler way for interpolation.
-            return Data(
-                x=torch.from_numpy(grp["x"][...]),
-                pos=torch.from_numpy(grp["pos"][...]),
-                y=torch.from_numpy(grp["y"][...]),
-                idx_in_original_cloud=grp["idx_in_original_cloud"][...],
-                # num_nodes=grp["pos"][...].shape[0],
-                x_features_names=grp["x"].attrs["x_features_names"].tolist(),
-            )
+        # Use property once to be sure that samples are all indexed into the hdf5 file.
+        self.samples_hdf5_paths
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx: int) -> Optional[Data]:
         sample_hdf5_path = self.samples_hdf5_paths[idx]
         data = self._get_data(sample_hdf5_path)
 
         # filter if empty
-        if self.pre_filter is not None and self.pre_filter(data):
+        if self.pre_filter and self.pre_filter(data):
             return None
 
         # Transforms, including sampling and some augmentations.
         transform = self.train_transform
-        if sample_hdf5_path.startswith("val") or sample_hdf5_path.startswith("test"):
+        if sample_hdf5_path.startswith("val") or sample_hdf5_path.startswith(
+            "test"
+        ):
             transform = self.eval_transform
-        if transform is not None:
+        if transform :
             data = transform(data)
 
         # filter if empty
-        if data is None or (self.pre_filter is not None and self.pre_filter(data)):
+        if not data or (
+            self.pre_filter and self.pre_filter(data)
+        ):
             return None
 
         return data
 
+    def _get_data(self, sample_hdf5_path: str) -> Data:
+        """Loads a Data object from the HDF5 dataset.
+
+        Opening the file has a high cost so we do it only once and store the opened files as a singleton
+        for each process within __get_item__ and not in __init__ to support for Multi-GPU.
+
+        See https://discuss.pytorch.org/t/dataloader-when-num-worker-0-there-is-bug/25643/16?u=piojanu.
+
+        """
+        if self.dataset is None:
+            self.dataset = h5py.File(self.hdf5_file_path, "r")
+
+        grp = self.dataset[sample_hdf5_path]
+        # [...] needed to make a copy of content and avoid closing HDF5.
+        # Nota: idx_in_original_cloud SHOULD be np.ndarray, in order to be batched into a list,
+        # which serves to keep track of indivual sample sizes in a simpler way for interpolation.
+        return Data(
+            x=torch.from_numpy(grp["x"][...]),
+            pos=torch.from_numpy(grp["pos"][...]),
+            y=torch.from_numpy(grp["y"][...]),
+            idx_in_original_cloud=grp["idx_in_original_cloud"][...],
+            x_features_names=grp["x"].attrs["x_features_names"].tolist(),
+            # num_nodes=grp["pos"][...].shape[0],  # Not needed - performed under the hood.
+        )
+
+    def __len__(self):
+        return len(self.samples_hdf5_paths)
+
     @property
     def traindata(self):
         return self._get_split_subset("train")
@@ -198,7 +157,114 @@ def testdata(self):
         return self._get_split_subset("test")
 
     def _get_split_subset(self, split: SPLIT_TYPE):
+        """Get a sub-dataset of a specific (train/val/test) split."""
         indices = [
-            idx for idx, p in enumerate(self.samples_hdf5_paths) if p.startswith(split)
+            idx
+            for idx, p in enumerate(self.samples_hdf5_paths)
+            if p.startswith(split)
         ]
         return torch.utils.data.Subset(self, indices)
+
+    def _add_to_dataset(self, split: str, las_path: str):
+        """Add samples from LAS into HDF5 dataset file.
+
+        Args:
+            split (str): specifies either "train", "val", or "test" split.
+            las_path (str): path to point cloud.
+
+        """
+        basename = osp.basename(las_path)
+
+        # Delete dataset for incomplete LAS entry, to start from scratch.
+        # Useful in case data preparation was interrupted.
+        with h5py.File(self.hdf5_file_path, "a") as hdf5_file:
+            if basename in hdf5_file[split] and "is_complete" not in hdf5_file[split][basename].attrs:
+                del hdf5_file[basename]
+                # Parse and add subtiles to split group.
+        with h5py.File(self.hdf5_file_path, "a") as hdf5_file:
+            if basename not in hdf5_file[split]:
+                sample_number = 0
+                for sample_number, (sample_idx, sample_points) in enumerate(split_cloud_into_samples(
+                    las_path,
+                    self.tile_width,
+                    self.subtile_width,
+                    self.subtile_shape,
+                    # No overlap at eval time.
+                    subtile_overlap=self.subtile_overlap_train
+                    if split == "train"
+                    else 0,
+                )):
+                    data = self.points_pre_transform(sample_points)
+                    if self.pre_filter is not None and self.pre_filter(data):
+                        # e.g. pre_filter spots situations where num_nodes is too small.
+                        continue
+                    sample_id = str(sample_number).zfill(5)
+                    hdf5_path = osp.join(split, basename, sample_id)
+                    hd5f_path_x = osp.join(hdf5_path, "x")
+                    hdf5_file.create_dataset(
+                        hd5f_path_x,
+                        data.x.shape,
+                        dtype="f",
+                        data=data.x,
+                    )
+                    hdf5_file[hd5f_path_x].attrs["x_features_names"] = copy.deepcopy(
+                        data.x_features_names
+                    )
+                    hdf5_file.create_dataset(
+                        osp.join(hdf5_path, "pos"),
+                        data.pos.shape,
+                        dtype="f",
+                        data=data.pos,
+                    )
+                    hdf5_file.create_dataset(
+                        osp.join(hdf5_path, "y"),
+                        data.y.shape,
+                        dtype="i",
+                        data=data.y,
+                    )
+                    hdf5_file.create_dataset(
+                        osp.join(hdf5_path, "idx_in_original_cloud"),
+                        sample_idx.shape,
+                        dtype="i",
+                        data=sample_idx,
+                    )
+
+                # A termination flag to report that all samples for this point cloud were included in the df5 file.
+                hdf5_file[split][basename].attrs["is_complete"] = True
+
+    @property
+    def samples_hdf5_paths(self):
+        """Index all samples in the dataset, if not already done before."""
+        # Use existing if already loaded as variable.
+        if self._samples_hdf5_paths:
+            return self._samples_hdf5_paths
+
+        # Load as variable if already indexed in hdf5 file. Need to decode b-string.
+        with h5py.File(self.hdf5_file_path, "r") as hdf5_file:
+            if "samples_hdf5_paths" in hdf5_file:
+                self._samples_hdf5_paths = [
+                    sample_path.decode("utf-8")
+                    for sample_path in hdf5_file["samples_hdf5_paths"]
+                ]
+                return self._samples_hdf5_paths
+
+        # Otherwise, index samples, and add the index as an attribute to the HDF5 file.
+        self._samples_hdf5_paths = []
+        with h5py.File(self.hdf5_file_path, "r") as hdf5_file:
+            for split in hdf5_file.keys():
+                if split not in ["train", "val", "test"]:
+                    continue
+                for basename in hdf5_file[split].keys():
+                    for sample_number in hdf5_file[split][basename].keys():
+                        self._samples_hdf5_paths.append(osp.join(split, basename, sample_number))
+
+        with h5py.File(self.hdf5_file_path, "a") as hdf5_file:
+            # special type to avoid silent string truncation in hdf5 datasets.
+            variable_lenght_str_datatype = h5py.special_dtype(vlen=str)
+            hdf5_file.create_dataset(
+                "samples_hdf5_paths",
+                (len(self.samples_hdf5_paths),),
+                dtype=variable_lenght_str_datatype,
+                data=self._samples_hdf5_paths,
+            )
+        return self._samples_hdf5_paths
diff --git a/myria3d/pctl/dataset/iterable.py b/myria3d/pctl/dataset/iterable.py
index a3b97863..5a9df515 100644
--- a/myria3d/pctl/dataset/iterable.py
+++ b/myria3d/pctl/dataset/iterable.py
@@ -13,13 +13,17 @@
 
 
 class InferenceDataset(IterableDataset):
-    """No need for HDF5 for Inference."""
+    """Iterable dataset to load samples from a single las file."""
 
     def __init__(
         self,
         las_file: str,
-        points_pre_transform: Callable[[ArrayLike], Data] = lidar_hd_pre_transform,
-        pre_filter: Optional[Callable[[Data], bool]] = pre_filter_below_n_points,
+        points_pre_transform: Callable[
+            [ArrayLike], Data
+        ] = lidar_hd_pre_transform,
+        pre_filter: Optional[
+            Callable[[Data], bool]
+        ] = pre_filter_below_n_points,
         transform: Optional[Callable[[Data], Data]] = None,
         tile_width: Number = 1000,
         subtile_width: Number = 50,
@@ -56,14 +60,14 @@ def get_iterator(self):
             # for final interpolation - should be kept as a np.ndarray to be batched as a list later.
             sample_data["idx_in_original_cloud"] = idx_in_original_cloud
 
-            if self.pre_filter is not None and self.pre_filter(sample_data):
+            if self.pre_filter and self.pre_filter(sample_data):
                 # e.g. not enough points in this receptive field.
                 continue
 
             if self.transform:
                 sample_data = self.transform(sample_data)
 
-            if self.pre_filter is not None and self.pre_filter(sample_data):
+            if self.pre_filter and self.pre_filter(sample_data):
                 # e.g. not enough points in this receptive field.
                 continue
 
diff --git a/myria3d/pctl/dataset/toy_dataset.py b/myria3d/pctl/dataset/toy_dataset.py
index 42276841..a7dbe9d4 100755
--- a/myria3d/pctl/dataset/toy_dataset.py
+++ b/myria3d/pctl/dataset/toy_dataset.py
@@ -9,8 +9,7 @@
 from myria3d.pctl.dataset.hdf5 import HDF5Dataset  # noqa
 
 
-TOY_LAS_DATA = "tests/data/toy_dataset_src/870000_6618000.subset.50mx100m.las"
-TOY_SPLIT_CSV = "tests/data/toy_dataset_src/toy_dataset_split.csv"
+TOY_LAS_DATA = "tests/data/toy_dataset_src/862000_6652000.classified_toy_dataset.100mx100m.las"
 TOY_DATASET_HDF5_PATH = "tests/data/toy_dataset.hdf5"
 
 
@@ -38,7 +37,7 @@ def make_toy_dataset_from_test_file():
     # TODO: update transforms ? or use a config ?
     HDF5Dataset(
         TOY_DATASET_HDF5_PATH,
-        las_files_by_split={
+        las_paths_by_split_dict={
             "train": [TOY_LAS_DATA],
             "val": [TOY_LAS_DATA],
             "test": [TOY_LAS_DATA],
diff --git a/myria3d/pctl/dataset/utils.py b/myria3d/pctl/dataset/utils.py
index 39a763ed..4ce8f010 100644
--- a/myria3d/pctl/dataset/utils.py
+++ b/myria3d/pctl/dataset/utils.py
@@ -10,7 +10,7 @@
 
 SPLIT_TYPE = Union[Literal["train"], Literal["val"], Literal["test"]]
 SHAPE_TYPE = Union[Literal["disk"], Literal["square"]]
-LAS_FILES_BY_SPLIT_TYPE = Dict[SPLIT_TYPE, List[str]]
+LAS_PATHS_BY_SPLIT_DICT_TYPE = Dict[SPLIT_TYPE, List[str]]
 
 # commons
 
@@ -119,7 +119,9 @@ def split_cloud_into_samples(
             # Adapt radius to have complete coverage of the data, with a slight overlap between samples.
             minkowski_p = 2
             radius = radius * math.sqrt(2)
-        sample_idx = np.array(kd_tree.query_ball_point(center, r=radius, p=minkowski_p))
+        sample_idx = np.array(
+            kd_tree.query_ball_point(center, r=radius, p=minkowski_p)
+        )
         if not len(sample_idx):
             # no points in this receptive fields
             continue
diff --git a/myria3d/pctl/transforms/transforms.py b/myria3d/pctl/transforms/transforms.py
index a4f108d2..7686ae6b 100755
--- a/myria3d/pctl/transforms/transforms.py
+++ b/myria3d/pctl/transforms/transforms.py
@@ -1,3 +1,5 @@
+import math
+import re
 from typing import Dict, List
 
 import numpy as np
@@ -22,6 +24,39 @@ def __call__(self, data: Data):
         return data
 
 
+class MinimumNumNodes(BaseTransform):
+    def __init__(self, num: int):
+        self.num = num
+
+    def __call__(self, data):
+        num_nodes = data.num_nodes
+
+        if num_nodes >= self.num:
+            return data
+
+        choice = torch.cat(
+            [torch.randperm(num_nodes) for _ in range(math.ceil(self.num / num_nodes))],
+            dim=0,
+        )[: self.num]
+
+        for key, item in data:
+            if key == "num_nodes":
+                data.num_nodes = choice.size(0)
+            elif bool(re.search("edge", key)):
+                continue
+            elif (
+                torch.is_tensor(item)
+                and item.size(0) == num_nodes
+                and item.size(0) != 1
+            ):
+                data[key] = item[choice]
+
+        return data
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.num}"
+
+
 class CopyFullPos:
     """Make a copy of the original positions - to be used for test and inference."""
 
diff --git a/myria3d/train.py b/myria3d/train.py
index 2f4b241e..3f1d4bb3 100755
--- a/myria3d/train.py
+++ b/myria3d/train.py
@@ -26,6 +26,8 @@
 
 log = utils.get_logger(__name__)
 
+NEURAL_NET_ARCHITECTURE_CONFIG_GROUP = "neural_net"
+
 
 def train(config: DictConfig) -> Trainer:
     """Training pipeline (+ Test, + Finetuning)
@@ -78,7 +80,7 @@ def train(config: DictConfig) -> Trainer:
     # Init lightning callbacks
     callbacks: List[Callback] = []
     if "callbacks" in config:
-        for _, cb_conf in config.callbacks.items():
+        for cb_conf in config.callbacks.values():
             if "_target_" in cb_conf:
                 log.info(f"Instantiating callback <{cb_conf._target_}>")
                 callbacks.append(hydra.utils.instantiate(cb_conf))
@@ -86,7 +88,7 @@ def train(config: DictConfig) -> Trainer:
     # Init lightning loggers
     logger: List[LightningLoggerBase] = []
     if "logger" in config:
-        for _, lg_conf in config.logger.items():
+        for lg_conf in config.logger.values():
             if "_target_" in lg_conf:
                 log.info(f"Instantiating logger <{lg_conf._target_}>")
                 logger.append(hydra.utils.instantiate(lg_conf))
@@ -154,8 +156,8 @@ def train(config: DictConfig) -> Trainer:
             )
             config.model.ckpt_path = trainer.checkpoint_callback.best_model_path
         log.info(
-                f"Test will use specified model checkpointed at \n {config.model.ckpt_path}"
-            )
+            f"Test will use specified model checkpointed at \n {config.model.ckpt_path}"
+        )
         trainer.test(
             model=model, datamodule=datamodule, ckpt_path=config.model.ckpt_path
         )
@@ -166,12 +168,7 @@ def train(config: DictConfig) -> Trainer:
         # Instantiates the Model but overwrites everything with current config,
         # except module related params (nnet architecture)
         kwargs_to_override = copy.deepcopy(model.hparams)
-        NEURAL_NET_ARCHITECTURE_CONFIG_GROUP = "neural_net"
-        kwargs_to_override = {
-            key: value
-            for key, value in kwargs_to_override.items()
-            if NEURAL_NET_ARCHITECTURE_CONFIG_GROUP not in key
-        }
+        kwargs_to_override.pop(NEURAL_NET_ARCHITECTURE_CONFIG_GROUP, None)  # removes that key if it's there
         model = Model.load_from_checkpoint(config.model.ckpt_path, **kwargs_to_override)
         trainer.fit(model=model, datamodule=datamodule, ckpt_path=None)
         log.info(f"Best checkpoint:\n{trainer.checkpoint_callback.best_model_path}")
diff --git a/package_metadata.yaml b/package_metadata.yaml
index 58726ff0..99763064 100644
--- a/package_metadata.yaml
+++ b/package_metadata.yaml
@@ -1,4 +1,4 @@
-__version__: "2.4.5"
+__version__: "3.0.2"
 __name__: "myria3d"
 __url__: "https://github.com/IGNF/myria3d"
 __description__: "Multiclass Semantic Segmentation for Lidar Point Cloud"
diff --git a/setup_env/requirements.txt b/setup_env/requirements.txt
deleted file mode 100755
index 6f768267..00000000
--- a/setup_env/requirements.txt
+++ /dev/null
@@ -1,16 +0,0 @@
-# --------- hydra --------- #
-hydra-core==1.1.*
-hydra-colorlog==1.1.*
-
-# --------- RST Linter --------- #
-rstcheck==3.3.*
-
-# --------- Documentation --------- #
-sphinx==4.5.*
-sphinx_rtd_theme==1.0.*
-myst_parser==0.17.*
-sphinx_paramlinks==0.5.*
-recommonmark==0.7.*
-sphinxnotes-mock==1.0.0b0  # still a beta
-sphinx-argparse==0.3.*  #  Using 
-docutils==0.17
\ No newline at end of file
diff --git a/setup_env/requirements.yml b/setup_env/requirements.yml
deleted file mode 100755
index 75cddfc1..00000000
--- a/setup_env/requirements.yml
+++ /dev/null
@@ -1,47 +0,0 @@
-name: myria3d
-channels:
-  - comet_ml
-  - pytorch
-  - pyg
-  - conda-forge
-  - anaconda
-dependencies:
-  - python==3.9.*
-  - pip
-  # --------- data formars --------- #
-  - numpy==1.20
-  - h5py
-  # --------- geo --------- #
-  - pygeos
-  - laspy
-  - shapely
-  - python-pdal  # useful for data preparation
-  # --------- deep learning --------- #
-  - pytorch==1.10.1
-  - pyg==2.0.3
-  - torchvision==0.11.2
-  - cudatoolkit=11.3  # single equal sign there, not a typo
-  - pytorch-lightning==1.5.9
-  - torchmetrics==0.7.*  # Else, pytorch-lightning will install the latest
-  # --------- loggers --------- #
-  - comet_ml
-  # --------- linters --------- #
-  - pre-commit # hooks for applying linters on commit
-  - black # code formatting
-  - isort # import sorting
-  - flake8 # code analysis
-  # --------- tests --------- #
-  - pytest==7.1.*
-  - coverage==6.3.*
-  - pytest-cov==3.0.*
-  # --------- others --------- #
-  - python-dotenv # loading env variables from .env file
-  - rich # beautiful text formatting in terminal
-  - sh # for running bash commands in some tests
-  - scikit-learn # used in some callbacks
-  - seaborn # used in some callbacks
-  - jupyterlab # better jupyter notebooks
-  - pudb # debugger
-  - numba==0.55.1
-  - pip:
-      - -r requirements.txt
\ No newline at end of file
diff --git a/setup_env/setup_env.sh b/setup_env/setup_env.sh
deleted file mode 100755
index 874255b1..00000000
--- a/setup_env/setup_env.sh
+++ /dev/null
@@ -1,9 +0,0 @@
-# WARNING: you may need to adapt cudatoolkit version to your needs.
-# Or remove it altogether in case you do not need CUDA support at all
-
-# Install mamba to create the environment faster
-conda install -y mamba -n base -c conda-forge
-# Build it
-mamba env create -f setup_env/requirements.yml
-# activate it
-conda activate myria3d
\ No newline at end of file


Test patch:
diff --git a/tests/data/toy_dataset.hdf5 b/tests/data/toy_dataset.hdf5
deleted file mode 100644
index 97365223..00000000
Binary files a/tests/data/toy_dataset.hdf5 and /dev/null differ
diff --git a/tests/data/toy_dataset_src/862000_6652000.classified_toy_dataset.100mx100m.las b/tests/data/toy_dataset_src/862000_6652000.classified_toy_dataset.100mx100m.las
new file mode 100644
index 00000000..187b1db7
Binary files /dev/null and b/tests/data/toy_dataset_src/862000_6652000.classified_toy_dataset.100mx100m.las differ
diff --git a/tests/data/toy_dataset_src/870000_6618000.subset.50mx100m.las b/tests/data/toy_dataset_src/870000_6618000.subset.50mx100m.las
deleted file mode 100644
index ef555276..00000000
Binary files a/tests/data/toy_dataset_src/870000_6618000.subset.50mx100m.las and /dev/null differ
diff --git a/tests/data/toy_dataset_src/toy_dataset_split.csv b/tests/data/toy_dataset_src/toy_dataset_split.csv
index b813dc07..36259314 100644
--- a/tests/data/toy_dataset_src/toy_dataset_split.csv
+++ b/tests/data/toy_dataset_src/toy_dataset_split.csv
@@ -1,7 +1,4 @@
 basename,split
-870000_6618000.subset.50mx100m.train1.las,train
-870000_6618000.subset.50mx100m.train2.las,train
-870000_6618000.subset.50mx100m.val1.las,val
-870000_6618000.subset.50mx100m.val2.las,val
-870000_6618000.subset.50mx100m.test1.las,test
-870000_6618000.subset.50mx100m.test2.las,test
\ No newline at end of file
+862000_6652000.classified_toy_dataset.100mx100m.las,train
+862000_6652000.classified_toy_dataset.100mx100m.las,val
+862000_6652000.classified_toy_dataset.100mx100m.las,test
\ No newline at end of file
diff --git a/tests/myria3d/models/modules/test_randla_net.py b/tests/myria3d/models/modules/test_randla_net.py
deleted file mode 100644
index ee7f6a32..00000000
--- a/tests/myria3d/models/modules/test_randla_net.py
+++ /dev/null
@@ -1,33 +0,0 @@
-import pytest
-import torch
-from torch_geometric.data import Batch
-from myria3d.models.modules.randla_net import RandLANet
-
-
-@pytest.mark.parametrize("num_graphs", [1, 4])
-def test_fake_run_randlanet(num_graphs):
-    """Documents expected data format and make a forward pass with RandLa-Net
-
-    Model pass with "batch_size=1" is a edge case that needs to pass to avoid unexpected crash due to incomplete batch at the
-    end of an inference.
-
-    """
-    num_euclidian_dimensions = 3
-    num_features = 9
-    d_in = num_euclidian_dimensions + num_features
-    num_classes = 6
-    hparams_net = {
-        "d_in": d_in,
-        "num_neighbors": 16,
-        "decimation": 4,
-        "dropout": 0.0,
-        "num_classes": num_classes,
-    }
-    batch = Batch()
-    batch.num_graphs = 1
-    num_points = 12500
-    batch.pos = torch.rand((num_points * batch.num_graphs, num_euclidian_dimensions))
-    batch.x = torch.rand((num_points * batch.num_graphs, num_features))
-    rln = RandLANet(hparams_net)
-    output = rln(batch)
-    assert output.shape == torch.Size([num_points * batch.num_graphs, num_classes])
diff --git a/tests/myria3d/models/modules/test_randla_nets.py b/tests/myria3d/models/modules/test_randla_nets.py
new file mode 100644
index 00000000..2231e6e7
--- /dev/null
+++ b/tests/myria3d/models/modules/test_randla_nets.py
@@ -0,0 +1,64 @@
+import pytest
+import torch
+from torch_geometric.data import Batch, Data
+from myria3d.models.modules.randla_net import RandLANet
+from myria3d.models.modules.pyg_randla_net import PyGRandLANet
+
+
+@pytest.mark.parametrize("num_graphs", [1, 2])
+def test_fake_run_randlanet(num_graphs):
+    """Documents expected data format and make a forward pass with RandLa-Net
+
+    Model pass with "batch_size=1" is a edge case that needs to pass to avoid unexpected crash due to incomplete batch at the
+    end of an inference.
+
+    """
+    num_euclidian_dimensions = 3
+    num_features = 9
+    d_in = num_euclidian_dimensions + num_features
+    num_classes = 6
+    hparams_net = {
+        "d_in": d_in,
+        "num_neighbors": 16,
+        "decimation": 4,
+        "dropout": 0.0,
+        "num_classes": num_classes,
+    }
+    batch = Batch()
+    batch.num_graphs = num_graphs
+    num_points = 12500
+    batch.pos = torch.rand((num_points * batch.num_graphs, num_euclidian_dimensions))
+    batch.x = torch.rand((num_points * batch.num_graphs, num_features))
+    rln = RandLANet(**hparams_net)
+    output = rln(batch)
+    assert output.shape == torch.Size([num_points * batch.num_graphs, num_classes])
+
+
+@pytest.mark.parametrize("num_nodes", [[12500, 12500], [50, 50], [12500, 10000]])
+def test_fake_run_pyg_randlanet(num_nodes):
+    """Documents expected data format and make a forward pass with PyG RandLa-Net.
+
+    Accepts small clouds even though decimation should lead to empty cloud.
+    Accepts point clouds of various sizes.
+
+    """
+    num_euclidian_dimensions = 3
+    num_features = 9
+    num_classes = 6
+    decimation = 4
+    num_neighbors = 16
+
+    batch = Batch.from_data_list(
+        [
+            Data(
+                x=torch.rand((n, num_features)),
+                pos=torch.rand((n, num_euclidian_dimensions)),
+                batch=torch.full((n,), idx),
+            )
+            for idx, n in enumerate(num_nodes)
+        ]
+    )
+
+    rln = PyGRandLANet(num_euclidian_dimensions + num_features, num_classes, decimation=decimation, num_neighbors=num_neighbors)
+    output = rln(batch)
+    assert output.shape == torch.Size([sum(num_nodes), num_classes])
diff --git a/tests/myria3d/test_train_and_predict.py b/tests/myria3d/test_train_and_predict.py
index 2ad4346b..1221f94a 100644
--- a/tests/myria3d/test_train_and_predict.py
+++ b/tests/myria3d/test_train_and_predict.py
@@ -243,7 +243,7 @@ def _make_list_of_necesary_hydra_overrides_with_tmp_paths(
     """Get list of overrides for hydra, the ones that are always needed when calling train/test.
 
     Args:
-        toy_dataset_hdf5_path (str): path to directory to dataset derived from large las
+        toy_dataset_hdf5_path (str): path to directory to dataset.
         tmpdir (str): path to temporary directory.
 
     """