Repository: DeepRegNet/DeepReg

Issue:
Log more metrics in tensorboard
## Subject of the feature

What are you trying to do and how would you want to do it differently?

Trying to debug using tensorboard, to see if any metric blows up.

Is it something you currently you cannot do?

Yes

Is this related to an issue/problem?

No

Has the feature been requested before? If yes, please provide a link to the issue.

No


![image](https://user-images.githubusercontent.com/4066929/111922853-3c83de00-8a94-11eb-8976-df4d9fc5739a.png)

current tensorboard contains limited metrics, we need to track more. In particular, the loss during the epoch should be tracked.


TODO
- [ ] Allow to adjust `update_freq`
- [ ] log TRE if data are labeled
- [ ] log vis of ddf, warped moving image, fixed image, etc.

Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 42becf05..d0089f17 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,7 @@ compatible with the updates.
 
 ### Added
 
+- Added statistics on inputs, DDF and TRE into tensorboard.
 - Added example for using custom loss.
 - Added tests on Mac OS.
 - Added tests for python 3.6 and 3.7.
diff --git a/deepreg/loss/label.py b/deepreg/loss/label.py
index 0d2c0b15..0e0d94ba 100644
--- a/deepreg/loss/label.py
+++ b/deepreg/loss/label.py
@@ -316,18 +316,16 @@ def compute_centroid(mask: tf.Tensor, grid: tf.Tensor) -> tf.Tensor:
     """
     Calculate the centroid of the mask.
     :param mask: shape = (batch, dim1, dim2, dim3)
-    :param grid: shape = (dim1, dim2, dim3, 3)
+    :param grid: shape = (1, dim1, dim2, dim3, 3)
     :return: shape = (batch, 3), batch of vectors denoting
              location of centroids.
     """
     assert len(mask.shape) == 4
-    assert len(grid.shape) == 4
+    assert len(grid.shape) == 5
     bool_mask = tf.expand_dims(
         tf.cast(mask >= 0.5, dtype=tf.float32), axis=4
     )  # (batch, dim1, dim2, dim3, 1)
-    masked_grid = bool_mask * tf.expand_dims(
-        grid, axis=0
-    )  # (batch, dim1, dim2, dim3, 3)
+    masked_grid = bool_mask * grid  # (batch, dim1, dim2, dim3, 3)
     numerator = tf.reduce_sum(masked_grid, axis=[1, 2, 3])  # (batch, 3)
     denominator = tf.reduce_sum(bool_mask, axis=[1, 2, 3])  # (batch, 1)
     return (numerator + EPS) / (denominator + EPS)  # (batch, 3)
@@ -340,7 +338,7 @@ def compute_centroid_distance(
     Calculate the L2-distance between two tensors' centroids.
     :param y_true: tensor, shape = (batch, dim1, dim2, dim3)
     :param y_pred: tensor, shape = (batch, dim1, dim2, dim3)
-    :param grid: tensor, shape = (dim1, dim2, dim3, 3)
+    :param grid: tensor, shape = (1, dim1, dim2, dim3, 3)
     :return: shape = (batch,)
     """
     centroid_1 = compute_centroid(mask=y_pred, grid=grid)  # (batch, 3)
diff --git a/deepreg/model/layer.py b/deepreg/model/layer.py
index d6b57bb7..fbeb6b08 100644
--- a/deepreg/model/layer.py
+++ b/deepreg/model/layer.py
@@ -244,8 +244,9 @@ class Warping(tfkl.Layer):
         super().__init__(name=name, **kwargs)
         self._fixed_image_size = fixed_image_size
         # shape = (1, f_dim1, f_dim2, f_dim3, 3)
-        self.grid_ref = layer_util.get_reference_grid(grid_size=fixed_image_size)
-        self.grid_ref = self.grid_ref[None, ...]
+        self.grid_ref = layer_util.get_reference_grid(grid_size=fixed_image_size)[
+            None, ...
+        ]
 
     def call(self, inputs, **kwargs) -> tf.Tensor:
         """
diff --git a/deepreg/model/network.py b/deepreg/model/network.py
index 8e764c09..227a20d4 100644
--- a/deepreg/model/network.py
+++ b/deepreg/model/network.py
@@ -6,7 +6,8 @@ from typing import Dict, Optional, Tuple
 
 import tensorflow as tf
 
-from deepreg.model import layer
+from deepreg.loss.label import DiceScore, compute_centroid_distance
+from deepreg.model import layer, layer_util
 from deepreg.model.backbone import GlobalNet
 from deepreg.registry import REGISTRY
 
@@ -64,6 +65,9 @@ class RegistrationModel(tf.keras.Model):
         self._inputs = None  # save inputs of self._model as dict
         self._outputs = None  # save outputs of self._model as dict
 
+        self.grid_ref = layer_util.get_reference_grid(grid_size=fixed_image_size)[
+            None, ...
+        ]
         self._model: tf.keras.Model = self.build_model()
         self.build_loss()
 
@@ -175,7 +179,7 @@ class RegistrationModel(tf.keras.Model):
         """
         Build and add one weighted loss together with the metrics.
 
-        :param name: name of loss
+        :param name: name of loss, image / label / regularization.
         :param inputs_dict: inputs for loss function
         """
 
@@ -233,6 +237,46 @@ class RegistrationModel(tf.keras.Model):
     def build_loss(self):
         """Build losses according to configs."""
 
+        # input metrics
+        fixed_image = self._inputs["fixed_image"]
+        moving_image = self._inputs["moving_image"]
+        self.log_tensor_stats(tensor=moving_image, name="moving_image")
+        self.log_tensor_stats(tensor=fixed_image, name="fixed_image")
+
+        # image loss, conditional model does not have this
+        if "pred_fixed_image" in self._outputs:
+            pred_fixed_image = self._outputs["pred_fixed_image"]
+            self._build_loss(
+                name="image",
+                inputs_dict=dict(y_true=fixed_image, y_pred=pred_fixed_image),
+            )
+
+        if self.labeled:
+            # input metrics
+            fixed_label = self._inputs["fixed_label"]
+            moving_label = self._inputs["moving_label"]
+            self.log_tensor_stats(tensor=moving_label, name="moving_label")
+            self.log_tensor_stats(tensor=fixed_label, name="fixed_label")
+
+            # label loss
+            pred_fixed_label = self._outputs["pred_fixed_label"]
+            self._build_loss(
+                name="label",
+                inputs_dict=dict(y_true=fixed_label, y_pred=pred_fixed_label),
+            )
+
+            # additional label metrics
+            tre = compute_centroid_distance(
+                y_true=fixed_label, y_pred=pred_fixed_label, grid=self.grid_ref
+            )
+            dice_binary = DiceScore(binary=True)(
+                y_true=fixed_label, y_pred=pred_fixed_label
+            )
+            self._model.add_metric(tre, name="metric/TRE", aggregation="mean")
+            self._model.add_metric(
+                dice_binary, name="metric/BinaryDiceScore", aggregation="mean"
+            )
+
     def call(
         self, inputs: Dict[str, tf.Tensor], training=None, mask=None
     ) -> Dict[str, tf.Tensor]:
@@ -287,6 +331,30 @@ class RegistrationModel(tf.keras.Model):
                 f"{err}"
             )
 
+    def log_tensor_stats(self, tensor: tf.Tensor, name: str):
+        """
+        Log statistics of a given tensor.
+
+        :param tensor: tensor to monitor.
+        :param name: name of the tensor.
+        """
+        flatten = tf.reshape(tensor, shape=(self.batch_size, -1))
+        self._model.add_metric(
+            tf.reduce_mean(flatten, axis=1),
+            name=f"metric/{name}_mean",
+            aggregation="mean",
+        )
+        self._model.add_metric(
+            tf.reduce_min(flatten, axis=1),
+            name=f"metric/{name}_min",
+            aggregation="min",
+        )
+        self._model.add_metric(
+            tf.reduce_max(flatten, axis=1),
+            name=f"metric/{name}_max",
+            aggregation="max",
+        )
+
 
 @REGISTRY.register_model(name="ddf")
 class DDFModel(RegistrationModel):
@@ -359,26 +427,12 @@ class DDFModel(RegistrationModel):
 
     def build_loss(self):
         """Build losses according to configs."""
-        fixed_image = self._inputs["fixed_image"]
-        ddf = self._outputs["ddf"]
-        pred_fixed_image = self._outputs["pred_fixed_image"]
+        super().build_loss()
 
-        # ddf
+        # ddf loss and metrics
+        ddf = self._outputs["ddf"]
         self._build_loss(name="regularization", inputs_dict=dict(inputs=ddf))
-
-        # image
-        self._build_loss(
-            name="image", inputs_dict=dict(y_true=fixed_image, y_pred=pred_fixed_image)
-        )
-
-        # label
-        if self.labeled:
-            fixed_label = self._inputs["fixed_label"]
-            pred_fixed_label = self._outputs["pred_fixed_label"]
-            self._build_loss(
-                name="label",
-                inputs_dict=dict(y_true=fixed_label, y_pred=pred_fixed_label),
-            )
+        self.log_tensor_stats(tensor=ddf, name="ddf")
 
     def postprocess(
         self,
@@ -456,9 +510,9 @@ class DVFModel(DDFModel):
         ddf = layer.IntDVF(fixed_image_size=self.fixed_image_size)(dvf)
 
         # build outputs
-        warping = layer.Warping(fixed_image_size=self.fixed_image_size)
+        self._warping = layer.Warping(fixed_image_size=self.fixed_image_size)
         # (f_dim1, f_dim2, f_dim3, 3)
-        pred_fixed_image = warping(inputs=[ddf, moving_image])
+        pred_fixed_image = self._warping(inputs=[ddf, moving_image])
 
         self._outputs = dict(dvf=dvf, ddf=ddf, pred_fixed_image=pred_fixed_image)
 
@@ -467,11 +521,19 @@ class DVFModel(DDFModel):
 
         # (f_dim1, f_dim2, f_dim3, 3)
         moving_label = self._inputs["moving_label"]
-        pred_fixed_label = warping(inputs=[ddf, moving_label])
+        pred_fixed_label = self._warping(inputs=[ddf, moving_label])
 
         self._outputs["pred_fixed_label"] = pred_fixed_label
         return tf.keras.Model(inputs=self._inputs, outputs=self._outputs)
 
+    def build_loss(self):
+        """Build losses according to configs."""
+        super().build_loss()
+
+        # dvf metrics
+        dvf = self._outputs["dvf"]
+        self.log_tensor_stats(tensor=dvf, name="dvf")
+
     def postprocess(
         self,
         inputs: Dict[str, tf.Tensor],
@@ -528,16 +590,6 @@ class ConditionalModel(RegistrationModel):
         self._outputs = dict(pred_fixed_label=pred_fixed_label)
         return tf.keras.Model(inputs=self._inputs, outputs=self._outputs)
 
-    def build_loss(self):
-        """Build losses according to configs."""
-        fixed_label = self._inputs["fixed_label"]
-        pred_fixed_label = self._outputs["pred_fixed_label"]
-
-        self._build_loss(
-            name="label",
-            inputs_dict=dict(y_true=fixed_label, y_pred=pred_fixed_label),
-        )
-
     def postprocess(
         self,
         inputs: Dict[str, tf.Tensor],
diff --git a/deepreg/train.py b/deepreg/train.py
index 5a318037..60aeabac 100644
--- a/deepreg/train.py
+++ b/deepreg/train.py
@@ -142,7 +142,9 @@ def train(
 
     # build callbacks
     tensorboard_callback = tf.keras.callbacks.TensorBoard(
-        log_dir=log_dir, histogram_freq=config["train"]["save_period"]
+        log_dir=log_dir,
+        histogram_freq=config["train"]["save_period"],
+        update_freq=config["train"].get("update_freq", "epoch"),
     )
     ckpt_callback, initial_epoch = build_checkpoint_callback(
         model=model,
diff --git a/deepreg/util.py b/deepreg/util.py
index dfad4508..fa37266e 100644
--- a/deepreg/util.py
+++ b/deepreg/util.py
@@ -171,7 +171,7 @@ def calculate_metrics(
         y_pred = pred_fixed_label[sample_index : (sample_index + 1), :, :, :]
         dice = label_loss.DiceScore(binary=True)(y_true=y_true, y_pred=y_pred).numpy()
         tre = label_loss.compute_centroid_distance(
-            y_true=y_true, y_pred=y_pred, grid=fixed_grid_ref[0, :, :, :, :]
+            y_true=y_true, y_pred=y_pred, grid=fixed_grid_ref
         ).numpy()[0]
     else:
         dice = None


Test patch:
diff --git a/test/unit/test_loss_label.py b/test/unit/test_loss_label.py
index 470bc35d..ba43f2ec 100644
--- a/test/unit/test_loss_label.py
+++ b/test/unit/test_loss_label.py
@@ -202,9 +202,9 @@ def test_compute_centroid():
     tensor_mask[0, :, :, :] = np.ones((2, 2, 2))
     tensor_mask = tf.constant(tensor_mask, dtype=tf.float32)
 
-    tensor_grid = np.ones((2, 2, 2, 3))
-    tensor_grid[:, :, :, 1] *= 2
-    tensor_grid[:, :, :, 2] *= 3
+    tensor_grid = np.ones((1, 2, 2, 2, 3))
+    tensor_grid[:, :, :, :, 1] *= 2
+    tensor_grid[:, :, :, :, 2] *= 3
     tensor_grid = tf.constant(tensor_grid, dtype=tf.float32)
 
     expected = np.ones((3, 3))  # use 1 because 0/0 ~= (0+eps)/(0+eps) = 1
@@ -223,8 +223,8 @@ def test_compute_centroid_d():
     tensor_mask[0, :, :, :] = array_ones
     tensor_mask = tf.convert_to_tensor(tensor_mask, dtype=tf.float32)
 
-    tensor_grid = np.zeros((2, 2, 2, 3))
-    tensor_grid[:, :, :, 0] = array_ones
+    tensor_grid = np.zeros((1, 2, 2, 2, 3))
+    tensor_grid[:, :, :, :, 0] = array_ones
     tensor_grid = tf.convert_to_tensor(tensor_grid, dtype=tf.float32)
 
     get = label.compute_centroid_distance(tensor_mask, tensor_mask, tensor_grid)