Repository: DeepRegNet/DeepReg

Issue:
Add docs for custom loss using registry
## Subject of the documentation

With the introduction of the registry, using custom loss/backbone/model should be easier. The related outdated docs have been removed. New docs will be necessary before the next release.

## TODO

- [x] backbone
- [ ] loss

Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 343fa7c7..4ac2973a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,7 @@ compatible with the updates.
 
 ### Added
 
+- Added example for using custom loss.
 - Added tests on Mac OS.
 - Added tests for python 3.6 and 3.7.
 - Added support to custom layer channels in U-Net.
diff --git a/docs/source/docs/registry.md b/docs/source/docs/registry.md
index 77713022..41f12f06 100644
--- a/docs/source/docs/registry.md
+++ b/docs/source/docs/registry.md
@@ -89,10 +89,12 @@ unet = REGISTRY.build_backbone(config=config)
 
 ## Example usages
 
-To further explain how to use `REGISTRY` for using customized classes, detailed examples
-are provided for the following categories:
+Apart from the table of [registered classes](registered_classes.html), to further
+explain how to use `REGISTRY` for using customized classes, detailed examples are
+provided for the following categories:
 
 - backbone
+- loss
 
 ### Custom backbone
 
@@ -101,8 +103,22 @@ To register a custom backbone class, the steps are as follows
 1. Subclass the `Backbone` and implement a custom backbone class.
 2. Import `REGISTRY` and use the decorator `@REGISTRY.register_backbone` to register the
    custom class.
-3. Use the registered name in the config for using the registered backbone.
+3. Use the registered name in the config for using the registered custom backbone.
 
 Please check the self-contained
 [example script](https://github.com/DeepRegNet/DeepReg/blob/main/examples/custom_backbone.py)
 for further details.
+
+### Custom Loss
+
+To register a custom loss class for images and labels, the steps are as follows
+
+1. Subclass the `tf.keras.losses.Loss` and implement a custom backbone class.
+2. Import `REGISTRY` and use the decorator `@REGISTRY.register_loss` to register the
+   custom class.
+3. Use the registered name in the config for using the registered custom loss.
+
+Please check the self-contained
+[example script](https://github.com/DeepRegNet/DeepReg/blob/main/examples/custom_image_label_loss.py)
+for further details. There is also a more complicated
+[example of parameterized custom loss](https://github.com/DeepRegNet/DeepReg/blob/main/examples/custom_parameterized_image_label_loss.py).
diff --git a/examples/config_custom_backbone.yaml b/examples/config_custom_backbone.yaml
index 8766ac63..834f65fd 100644
--- a/examples/config_custom_backbone.yaml
+++ b/examples/config_custom_backbone.yaml
@@ -30,9 +30,8 @@ train:
 
   # define the optimizer
   optimizer:
-    name: "adam" # options include "adam", "sgd" and "rms"
-    adam:
-      learning_rate: 1.0e-5
+    name: "Adam"
+    learning_rate: 1.0e-5
 
   # define the hyper-parameters for preprocessing
   preprocess:
diff --git a/examples/config_custom_image_label_loss.yaml b/examples/config_custom_image_label_loss.yaml
new file mode 100644
index 00000000..47b4f413
--- /dev/null
+++ b/examples/config_custom_image_label_loss.yaml
@@ -0,0 +1,45 @@
+dataset:
+  dir:
+    train: "data/test/nifti/unpaired/train" # required
+    valid: "data/test/nifti/unpaired/test" # required
+    test: "data/test/nifti/unpaired/test" # required
+  format: "nifti"
+  type: "unpaired" # paired / unpaired / grouped
+  labeled: true # whether to use the labels if available, "true" or "false"
+  image_shape: [16, 16, 16]
+
+train:
+  # define neural network structure
+  method: "ddf" # options include "ddf", "dvf", "conditional"
+  backbone:
+    name: "local" # options include "local", "unet" and "global"
+    num_channel_initial: 1 # number of initial channel in local net, controls the size of the network
+    extract_levels: [0, 1, 2]
+  # define the loss function for training
+  loss:
+    image:
+      name: "root_mean_square" # other options include "lncc", "ssd" and "gmi", for local normalised cross correlation,
+      weight: 0.1
+    label:
+      weight: 1.0
+      name: "dice" # options include "dice", "cross-entropy", "mean-squared", "generalised_dice" and "jaccard"
+      scales: [0, 1, 2, 4, 8, 16, 32]
+    regularization:
+      weight: 0.5 # weight of regularization loss
+      name: "bending" # options include "bending", "gradient"
+
+  # define the optimizer
+  optimizer:
+    name: "Adam"
+    learning_rate: 1.0e-5
+
+  # define the hyper-parameters for preprocessing
+  preprocess:
+    data_augmentation:
+      name: "affine"
+    batch_size: 2
+    shuffle_buffer_num_batch: 1 # shuffle_buffer_size = batch_size * shuffle_buffer_num_batch
+
+  # other training hyper-parameters
+  epochs: 2 # number of training epochs
+  save_period: 2 # the model will be saved every `save_period` epochs.
diff --git a/examples/config_custom_parameterized_image_label_loss.yaml b/examples/config_custom_parameterized_image_label_loss.yaml
new file mode 100644
index 00000000..ff72eab2
--- /dev/null
+++ b/examples/config_custom_parameterized_image_label_loss.yaml
@@ -0,0 +1,46 @@
+dataset:
+  dir:
+    train: "data/test/nifti/unpaired/train" # required
+    valid: "data/test/nifti/unpaired/test" # required
+    test: "data/test/nifti/unpaired/test" # required
+  format: "nifti"
+  type: "unpaired" # paired / unpaired / grouped
+  labeled: true # whether to use the labels if available, "true" or "false"
+  image_shape: [16, 16, 16]
+
+train:
+  # define neural network structure
+  method: "ddf" # options include "ddf", "dvf", "conditional"
+  backbone:
+    name: "local" # options include "local", "unet" and "global"
+    num_channel_initial: 1 # number of initial channel in local net, controls the size of the network
+    extract_levels: [0, 1, 2]
+  # define the loss function for training
+  loss:
+    image:
+      name: "lp_norm" # other options include "lncc", "ssd" and "gmi", for local normalised cross correlation,
+      p: 1 # 1 or 2
+      weight: 0.1
+    label:
+      weight: 1.0
+      name: "dice" # options include "dice", "cross-entropy", "mean-squared", "generalised_dice" and "jaccard"
+      scales: [0, 1, 2, 4, 8, 16, 32]
+    regularization:
+      weight: 0.5 # weight of regularization loss
+      name: "bending" # options include "bending", "gradient"
+
+  # define the optimizer
+  optimizer:
+    name: "Adam"
+    learning_rate: 1.0e-5
+
+  # define the hyper-parameters for preprocessing
+  preprocess:
+    data_augmentation:
+      name: "affine"
+    batch_size: 2
+    shuffle_buffer_num_batch: 1 # shuffle_buffer_size = batch_size * shuffle_buffer_num_batch
+
+  # other training hyper-parameters
+  epochs: 2 # number of training epochs
+  save_period: 2 # the model will be saved every `save_period` epochs.
diff --git a/examples/custom_image_label_loss.py b/examples/custom_image_label_loss.py
new file mode 100644
index 00000000..341a9e0a
--- /dev/null
+++ b/examples/custom_image_label_loss.py
@@ -0,0 +1,51 @@
+"""This script provides an example of using custom backbone for training."""
+import tensorflow as tf
+
+from deepreg.registry import REGISTRY
+from deepreg.train import train
+
+
+@REGISTRY.register_loss(name="root_mean_square")
+class RootMeanSquaredDifference(tf.keras.losses.Loss):
+    """
+    Square root of the mean of squared distance between y_true and y_pred.
+
+    y_true and y_pred have to be at least 1d tensor, including batch axis.
+    """
+
+    def __init__(
+        self,
+        reduction: str = tf.keras.losses.Reduction.SUM,
+        name: str = "RootMeanSquaredDifference",
+    ):
+        """
+        Init.
+
+        :param reduction: using SUM reduction over batch axis,
+            calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
+        :param name: name of the loss
+        """
+        super().__init__(reduction=reduction, name=name)
+
+    def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
+        """
+        Return loss for a batch.
+
+        :param y_true: shape = (batch, ...)
+        :param y_pred: shape = (batch, ...)
+        :return: shape = (batch,)
+        """
+        loss = tf.math.squared_difference(y_true, y_pred)
+        loss = tf.keras.layers.Flatten()(loss)
+        loss = tf.reduce_mean(loss, axis=1)
+        loss = tf.math.sqrt(loss)
+        return loss
+
+
+config_path = "examples/config_custom_image_label_loss.yaml"
+train(
+    gpu="",
+    config_path=config_path,
+    gpu_allow_growth=True,
+    ckpt_path="",
+)
diff --git a/examples/custom_parameterized_image_label_loss.py b/examples/custom_parameterized_image_label_loss.py
new file mode 100644
index 00000000..44010fac
--- /dev/null
+++ b/examples/custom_parameterized_image_label_loss.py
@@ -0,0 +1,55 @@
+"""This script provides an example of using custom backbone for training."""
+import tensorflow as tf
+
+from deepreg.registry import REGISTRY
+from deepreg.train import train
+
+
+@REGISTRY.register_loss(name="lp_norm")
+class LPNorm(tf.keras.losses.Loss):
+    """
+    L^p norm between y_true and y_pred, p = 1 or 2.
+
+    y_true and y_pred have to be at least 1d tensor, including batch axis.
+    """
+
+    def __init__(
+        self,
+        p: int,
+        reduction: str = tf.keras.losses.Reduction.SUM,
+        name: str = "LPNorm",
+    ):
+        """
+        Init.
+
+        :param p: order of the norm, 1 or 2.
+        :param reduction: using SUM reduction over batch axis,
+            calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
+        :param name: name of the loss.
+        """
+        super().__init__(reduction=reduction, name=name)
+        if p not in [1, 2]:
+            raise ValueError(f"For LPNorm, p must be 0 or 1, got {p}.")
+        self.p = p
+
+    def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
+        """
+        Return loss for a batch.
+
+        :param y_true: shape = (batch, ...)
+        :param y_pred: shape = (batch, ...)
+        :return: shape = (batch,)
+        """
+        diff = y_true - y_pred
+        diff = tf.keras.layers.Flatten()(diff)
+        loss = tf.norm(diff, axis=-1, ord=self.p)
+        return loss
+
+
+config_path = "examples/config_custom_parameterized_image_label_loss.yaml"
+train(
+    gpu="",
+    config_path=config_path,
+    gpu_allow_growth=True,
+    ckpt_path="",
+)


Test patch:
diff --git a/test/integration/test_examples.py b/test/integration/test_examples.py
index 21b09470..b819e29e 100644
--- a/test/integration/test_examples.py
+++ b/test/integration/test_examples.py
@@ -20,6 +20,8 @@ class TestTutorial:
         "name",
         [
             "custom_backbone",
+            "custom_image_label_loss",
+            "custom_parameterized_image_label_loss",
         ],
     )
     def test_registry(self, name: str):