Repository: DeepRegNet/DeepReg

Issue:
Nan/inf loss encountered
## Subject of the issue

More discussion is here https://github.com/Project-MONAI/MONAI/issues/1868

Using script https://github.com/DeepRegNet/Benchmark/blob/9-data-for-miccai/benchmark/balakrishnan2019/voxel_morph_balakrishnan_2019.py with the latest `main` branch.

```
226/226 [==============================] - 3083s 14s/step - loss: -0.7058 - loss/regularization_GradientNorm: 2.1987e-04 - loss/regularization_GradientNorm_weighted: 2.1987e-04 - loss/image_LocalNormalizedCrossCorrelationLoss: -0.7060 - loss/image_LocalNormalizedCrossCorrelationLoss_weighted: -0.7060
Epoch 2/2000
226/226 [==============================] - 3266s 14s/step - loss: -0.4778 - loss/regularization_GradientNorm: 0.3009 - loss/regularization_GradientNorm_weighted: 0.3009 - loss/image_LocalNormalizedCrossCorrelationLoss: -0.7787 - loss/image_LocalNormalizedCrossCorrelationLoss_weighted: -0.7787
Epoch 3/2000
226/226 [==============================] - 3280s 15s/step - loss: nan - loss/regularization_GradientNorm: nan - loss/regularization_GradientNorm_weighted: nan - loss/image_LocalNormalizedCrossCorrelationLoss: -inf - loss/image_LocalNormalizedCrossCorrelationLoss_weighted: -inf
```


## Your environment

- DeepReg version (commit hash)

  `5e5bbdc`

- OS (e.g. Ubuntu 20.04, MacOS 10.15, etc.)

  Ubuntu

- Python Version (3.7, 3.8, etc.)

  3.7

- TensorFlow

  2.3



Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d0089f17..720ca54f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -24,6 +24,7 @@ compatible with the updates.
 
 ### Changed
 
+- Increased all EPS to 1e-5.
 - Clarify the suggestion in doc to use all-zero masks for missing labels.
 - Moved contributor list to a separate page.
 - Changed `no-test` flag to `full` for demo scripts.
@@ -37,6 +38,7 @@ compatible with the updates.
 
 ### Fixed
 
+- Fixed LNCC loss regarding INF values.
 - Removed loss weight checks to be more robust.
 - Fixed import error under python 3.6.
 - Fixed the residual module in local net architecture, compatible for previous
diff --git a/deepreg/constant.py b/deepreg/constant.py
new file mode 100644
index 00000000..72746ee3
--- /dev/null
+++ b/deepreg/constant.py
@@ -0,0 +1,3 @@
+"""Module defining global constants."""
+
+EPS = 1.0e-5
diff --git a/deepreg/dataset/loader/interface.py b/deepreg/dataset/loader/interface.py
index 1c89d5e5..99744819 100644
--- a/deepreg/dataset/loader/interface.py
+++ b/deepreg/dataset/loader/interface.py
@@ -379,10 +379,10 @@ class GeneratorDataLoader(DataLoader, ABC):
         for arr, name in zip(
             [moving_image, fixed_image], ["moving_image", "fixed_image"]
         ):
-            if len(arr.shape) != 3:
+            if len(arr.shape) != 3 or min(arr.shape) <= 0:
                 raise ValueError(
-                    f"Sample {image_indices}'s {name}' shape should be 3D. "
-                    f"Got {arr.shape}."
+                    f"Sample {image_indices}'s {name}' shape should be 3D"
+                    f" and non-empty, got {arr.shape}."
                 )
         # when data are labeled
         if moving_label is not None and fixed_label is not None:
diff --git a/deepreg/loss/image.py b/deepreg/loss/image.py
index 44510bdd..fed72b05 100644
--- a/deepreg/loss/image.py
+++ b/deepreg/loss/image.py
@@ -1,6 +1,7 @@
 """Provide different loss or metrics classes for images."""
 import tensorflow as tf
 
+from deepreg.constant import EPS
 from deepreg.loss.util import NegativeLossMixin
 from deepreg.loss.util import gaussian_kernel1d_size as gaussian_kernel1d
 from deepreg.loss.util import (
@@ -10,8 +11,6 @@ from deepreg.loss.util import (
 )
 from deepreg.registry import REGISTRY
 
-EPS = tf.keras.backend.epsilon()
-
 
 @REGISTRY.register_loss(name="ssd")
 class SumSquaredDifference(tf.keras.losses.Loss):
@@ -30,6 +29,8 @@ class SumSquaredDifference(tf.keras.losses.Loss):
         Init.
 
         :param reduction: using SUM reduction over batch axis,
+            this is for supporting multi-device training,
+            and the loss will be divided by global batch size,
             calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
         :param name: name of the loss
         """
@@ -71,6 +72,8 @@ class GlobalMutualInformation(tf.keras.losses.Loss):
         :param num_bins: number of bins for intensity, the default value is empirical.
         :param sigma_ratio: a hyper param for gaussian function
         :param reduction: using SUM reduction over batch axis,
+            this is for supporting multi-device training,
+            and the loss will be divided by global batch size,
             calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
         :param name: name of the loss
         """
@@ -189,6 +192,8 @@ class LocalNormalizedCrossCorrelation(tf.keras.losses.Loss):
         self,
         kernel_size: int = 9,
         kernel_type: str = "rectangular",
+        smooth_nr: float = EPS,
+        smooth_dr: float = EPS,
         reduction: str = tf.keras.losses.Reduction.SUM,
         name: str = "LocalNormalizedCrossCorrelation",
     ):
@@ -197,7 +202,11 @@ class LocalNormalizedCrossCorrelation(tf.keras.losses.Loss):
 
         :param kernel_size: int. Kernel size or kernel sigma for kernel_type='gauss'.
         :param kernel_type: str, rectangular, triangular or gaussian
+        :param smooth_nr: small constant added to numerator in case of zero covariance.
+        :param smooth_dr: small constant added to denominator in case of zero variance.
         :param reduction: using SUM reduction over batch axis,
+            this is for supporting multi-device training,
+            and the loss will be divided by global batch size,
             calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
         :param name: name of the loss
         """
@@ -210,6 +219,8 @@ class LocalNormalizedCrossCorrelation(tf.keras.losses.Loss):
         self.kernel_fn = self.kernel_fn_dict[kernel_type]
         self.kernel_type = kernel_type
         self.kernel_size = kernel_size
+        self.smooth_nr = smooth_nr
+        self.smooth_dr = smooth_dr
 
         # (kernel_size, )
         self.kernel = self.kernel_fn(kernel_size=self.kernel_size)
@@ -220,24 +231,25 @@ class LocalNormalizedCrossCorrelation(tf.keras.losses.Loss):
             * self.kernel[None, None, :]
         )
 
-    def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
+    def calc_ncc(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
         """
-        Return loss for a batch.
-
-        :param y_true: shape = (batch, dim1, dim2, dim3)
-            or (batch, dim1, dim2, dim3, ch)
-        :param y_pred: shape = (batch, dim1, dim2, dim3)
-            or (batch, dim1, dim2, dim3, ch)
-        :return: shape = (batch,)
+        Return NCC for a batch.
+
+        The kernel should not be normalized, as normalizing them leads to computation
+        with small values and the precision will be reduced.
+        Here both numerator and denominator are actually multiplied by kernel volume,
+        which helps the precision as well.
+        However, when the variance is zero, the obtained value might be negative due to
+        machine error. Therefore a hard-coded clipping is added to
+        prevent division by zero.
+
+        :param y_true: shape = (batch, dim1, dim2, dim3, 1)
+        :param y_pred: shape = (batch, dim1, dim2, dim3, 1)
+        :return: shape = (batch, dim1, dim2, dim3. 1)
         """
-        # adjust
-        if len(y_true.shape) == 4:
-            y_true = tf.expand_dims(y_true, axis=4)
-            y_pred = tf.expand_dims(y_pred, axis=4)
-        assert len(y_true.shape) == len(y_pred.shape) == 5
 
         # t = y_true, p = y_pred
-        # (batch, dim1, dim2, dim3, ch)
+        # (batch, dim1, dim2, dim3, 1)
         t2 = y_true * y_true
         p2 = y_pred * y_pred
         tp = y_true * y_pred
@@ -260,16 +272,53 @@ class LocalNormalizedCrossCorrelation(tf.keras.losses.Loss):
         t_var = t2_sum - t_avg * t_sum  # V[t] * E[1]
         p_var = p2_sum - p_avg * p_sum  # V[p] * E[1]
 
+        # ensure variance >= 0
+        t_var = tf.maximum(t_var, 0)
+        p_var = tf.maximum(p_var, 0)
+
         # (E[tp] - E[p] * E[t]) ** 2 / V[t] / V[p]
-        ncc = (cross * cross + EPS) / (t_var * p_var + EPS)
+        ncc = (cross * cross + self.smooth_nr) / (t_var * p_var + self.smooth_dr)
 
+        return ncc
+
+    def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
+        """
+        Return loss for a batch.
+
+        TODO: support channel axis dimension > 1.
+
+        :param y_true: shape = (batch, dim1, dim2, dim3)
+            or (batch, dim1, dim2, dim3, 1)
+        :param y_pred: shape = (batch, dim1, dim2, dim3)
+            or (batch, dim1, dim2, dim3, 1)
+        :return: shape = (batch,)
+        """
+        # sanity checks
+        if len(y_true.shape) == 4:
+            y_true = tf.expand_dims(y_true, axis=4)
+        if y_true.shape[4] != 1:
+            raise ValueError(
+                "Last dimension of y_true is not one. " f"y_true.shape = {y_true.shape}"
+            )
+        if len(y_pred.shape) == 4:
+            y_pred = tf.expand_dims(y_pred, axis=4)
+        if y_pred.shape[4] != 1:
+            raise ValueError(
+                "Last dimension of y_pred is not one. " f"y_pred.shape = {y_pred.shape}"
+            )
+
+        ncc = self.calc_ncc(y_true=y_true, y_pred=y_pred)
         return tf.reduce_mean(ncc, axis=[1, 2, 3, 4])
 
     def get_config(self) -> dict:
         """Return the config dictionary for recreating this class."""
         config = super().get_config()
-        config["kernel_size"] = self.kernel_size
-        config["kernel_type"] = self.kernel_type
+        config.update(
+            kernel_size=self.kernel_size,
+            kernel_type=self.kernel_type,
+            smooth_nr=self.smooth_nr,
+            smooth_dr=self.smooth_dr,
+        )
         return config
 
 
diff --git a/deepreg/loss/label.py b/deepreg/loss/label.py
index 0e0d94ba..c5162161 100644
--- a/deepreg/loss/label.py
+++ b/deepreg/loss/label.py
@@ -4,13 +4,12 @@ from typing import List, Optional
 
 import tensorflow as tf
 
+from deepreg.constant import EPS
 from deepreg.loss.util import NegativeLossMixin, cauchy_kernel1d
 from deepreg.loss.util import gaussian_kernel1d_sigma as gaussian_kernel1d
 from deepreg.loss.util import separable_filter
 from deepreg.registry import REGISTRY
 
-EPS = tf.keras.backend.epsilon()
-
 
 class MultiScaleLoss(tf.keras.losses.Loss):
     """
@@ -35,6 +34,8 @@ class MultiScaleLoss(tf.keras.losses.Loss):
         :param scales: list of scalars or None, if None, do not apply any scaling.
         :param kernel: gaussian or cauchy.
         :param reduction: using SUM reduction over batch axis,
+            this is for supporting multi-device training,
+            and the loss will be divided by global batch size,
             calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
         :param name: str, name of the loss.
         """
@@ -133,6 +134,8 @@ class DiceScore(MultiScaleLoss):
         :param scales: list of scalars or None, if None, do not apply any scaling.
         :param kernel: gaussian or cauchy.
         :param reduction: using SUM reduction over batch axis,
+            this is for supporting multi-device training,
+            and the loss will be divided by global batch size,
             calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
         :param name: str, name of the loss.
         """
@@ -207,6 +210,8 @@ class CrossEntropy(MultiScaleLoss):
         :param scales: list of scalars or None, if None, do not apply any scaling.
         :param kernel: gaussian or cauchy.
         :param reduction: using SUM reduction over batch axis,
+            this is for supporting multi-device training,
+            and the loss will be divided by global batch size,
             calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
         :param name: str, name of the loss.
         """
@@ -273,6 +278,8 @@ class JaccardIndex(MultiScaleLoss):
         :param scales: list of scalars or None, if None, do not apply any scaling.
         :param kernel: gaussian or cauchy.
         :param reduction: using SUM reduction over batch axis,
+            this is for supporting multi-device training,
+            and the loss will be divided by global batch size,
             calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
         :param name: str, name of the loss.
         """
diff --git a/deepreg/loss/util.py b/deepreg/loss/util.py
index 6f453d73..05083461 100644
--- a/deepreg/loss/util.py
+++ b/deepreg/loss/util.py
@@ -27,13 +27,9 @@ class NegativeLossMixin(tf.keras.losses.Loss):
         return -super().call(y_true=y_true, y_pred=y_pred)
 
 
-EPS = tf.keras.backend.epsilon()
-
-
 def rectangular_kernel1d(kernel_size: int) -> tf.Tensor:
     """
-    Return a the 1D filter for separable convolution equivalent to a 3-D rectangular
-    kernel for LocalNormalizedCrossCorrelation.
+    Return a the 1D rectangular kernel for LocalNormalizedCrossCorrelation.
 
     :param kernel_size: scalar, size of the 1-D kernel
     :return: kernel_weights, of shape (kernel_size, )
@@ -45,7 +41,7 @@ def rectangular_kernel1d(kernel_size: int) -> tf.Tensor:
 
 def triangular_kernel1d(kernel_size: int) -> tf.Tensor:
     """
-    1D triangular kernel.
+    Return a the 1D triangular kernel for LocalNormalizedCrossCorrelation.
 
     Assume kernel_size is odd, it will be a smoothed from
     a kernel which center part is zero.
@@ -73,13 +69,14 @@ def triangular_kernel1d(kernel_size: int) -> tf.Tensor:
     kernel = tf.nn.conv1d(
         kernel[None, :, None], filters=filters, stride=[1, 1, 1], padding="SAME"
     )
+
     return kernel[0, :, 0]
 
 
 def gaussian_kernel1d_size(kernel_size: int) -> tf.Tensor:
     """
-    Return a the 1D filter for separable convolution equivalent to a 3-D Gaussian
-    kernel for LocalNormalizedCrossCorrelation.
+    Return a the 1D Gaussian kernel for LocalNormalizedCrossCorrelation.
+
     :param kernel_size: scalar, size of the 1-D kernel
     :return: filters, of shape (kernel_size, )
     """
diff --git a/deepreg/model/network.py b/deepreg/model/network.py
index 227a20d4..de1140d0 100644
--- a/deepreg/model/network.py
+++ b/deepreg/model/network.py
@@ -30,8 +30,8 @@ class RegistrationModel(tf.keras.Model):
 
     def __init__(
         self,
-        moving_image_size: tuple,
-        fixed_image_size: tuple,
+        moving_image_size: Tuple,
+        fixed_image_size: Tuple,
         index_size: int,
         labeled: bool,
         batch_size: int,
diff --git a/examples/custom_image_label_loss.py b/examples/custom_image_label_loss.py
index 341a9e0a..3dd29315 100644
--- a/examples/custom_image_label_loss.py
+++ b/examples/custom_image_label_loss.py
@@ -22,6 +22,8 @@ class RootMeanSquaredDifference(tf.keras.losses.Loss):
         Init.
 
         :param reduction: using SUM reduction over batch axis,
+            this is for supporting multi-device training,
+            and the loss will be divided by global batch size,
             calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
         :param name: name of the loss
         """
diff --git a/examples/custom_parameterized_image_label_loss.py b/examples/custom_parameterized_image_label_loss.py
index 44010fac..c0db057d 100644
--- a/examples/custom_parameterized_image_label_loss.py
+++ b/examples/custom_parameterized_image_label_loss.py
@@ -24,6 +24,8 @@ class LPNorm(tf.keras.losses.Loss):
 
         :param p: order of the norm, 1 or 2.
         :param reduction: using SUM reduction over batch axis,
+            this is for supporting multi-device training,
+            and the loss will be divided by global batch size,
             calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
         :param name: name of the loss.
         """


Test patch:
diff --git a/test/unit/test_interface.py b/test/unit/test_interface.py
index 0315938d..22778f61 100644
--- a/test/unit/test_interface.py
+++ b/test/unit/test_interface.py
@@ -356,7 +356,7 @@ def test_generator_data_loader(caplog):
             fixed_label=None,
             image_indices=[1],
         )
-    assert "Sample [1]'s moving_image' shape should be 3D. " in str(err_info.value)
+    assert "Sample [1]'s moving_image' shape should be 3D" in str(err_info.value)
     with pytest.raises(ValueError) as err_info:
         generator.validate_images_and_labels(
             fixed_image=dummy_array,
diff --git a/test/unit/test_loss_image.py b/test/unit/test_loss_image.py
index 9d90baf9..19330ee6 100644
--- a/test/unit/test_loss_image.py
+++ b/test/unit/test_loss_image.py
@@ -5,12 +5,14 @@ Notes: The format of inputs to the function dissimilarity_fn
 in image.py should be better converted into tf tensor type beforehand.
 """
 from test.unit.util import is_equal_tf
+from typing import Tuple
 
 import numpy as np
 import pytest
 import tensorflow as tf
 
 import deepreg.loss.image as image
+from deepreg.constant import EPS
 
 
 class TestSumSquaredDistance:
@@ -76,38 +78,160 @@ def test_kernel_fn(kernel_size, name):
 
 class TestLocalNormalizedCrossCorrelation:
     @pytest.mark.parametrize(
-        "y_true,y_pred,shape,kernel_type,expected",
+        ("y_true_shape", "y_pred_shape"),
         [
-            (0.6, 0.3, (12, 12, 12, 12), "rectangular", 1.0),
-            (0.6, 0.3, (12, 12, 12, 12, 1), "rectangular", 1.0),
-            (0.0, 1.0, (12, 12, 12, 12, 1), "rectangular", 1.0),
-            (0.6, 0.3, (12, 12, 12, 12, 1), "gaussian", 1.0),
-            (0.6, 0.3, (12, 12, 12, 12, 1), "triangular", 1.0),
+            ((2, 3, 4, 5), (2, 3, 4, 5)),
+            ((2, 3, 4, 5), (2, 3, 4, 5, 1)),
+            ((2, 3, 4, 5, 1), (2, 3, 4, 5)),
+            ((2, 3, 4, 5, 1), (2, 3, 4, 5, 1)),
         ],
     )
-    def test_zero_info(self, y_true, y_pred, shape, kernel_type, expected):
-        y_true = y_true * tf.ones(shape=shape)
-        y_pred = y_pred * tf.ones(shape=shape)
-        expected = expected * tf.ones(shape=(shape[0],))
-        got = image.LocalNormalizedCrossCorrelation(kernel_type=kernel_type).call(
+    def test_input_shape(self, y_true_shape: Tuple, y_pred_shape: Tuple):
+        """
+        Test input with / without channel axis.
+
+        :param y_true_shape: input shape for y_true.
+        :param y_pred_shape: input shape for y_pred.
+        """
+        y_true = tf.ones(shape=y_true_shape)
+        y_pred = tf.ones(shape=y_pred_shape)
+        got = image.LocalNormalizedCrossCorrelation().call(
+            y_true,
+            y_pred,
+        )
+        assert got.shape == y_true_shape[:1]
+
+    @pytest.mark.parametrize(
+        ("y_true_shape", "y_pred_shape", "name"),
+        [
+            ((2, 3, 4, 5), (2, 3, 4, 5, 6), "y_pred"),
+            ((2, 3, 4, 5, 6), (2, 3, 4, 5), "y_true"),
+        ],
+    )
+    def test_input_shape_err(self, y_true_shape: Tuple, y_pred_shape: Tuple, name: str):
+        """
+        Current LNCC does not support image having channel dimension > 1.
+
+        :param y_true_shape: input shape for y_true.
+        :param y_pred_shape: input shape for y_pred.
+        :param name: name of the tensor having error.
+        """
+        y_true = tf.ones(shape=y_true_shape)
+        y_pred = tf.ones(shape=y_pred_shape)
+        with pytest.raises(ValueError) as err_info:
+            image.LocalNormalizedCrossCorrelation().call(y_true, y_pred)
+        assert f"Last dimension of {name} is not one." in str(err_info.value)
+
+    @pytest.mark.parametrize("value", [0.0, 0.5, 1.0])
+    @pytest.mark.parametrize(
+        ("smooth_nr", "smooth_dr", "expected"),
+        [
+            (1e-5, 1e-5, 1),
+            (0, 1e-5, 0),
+            (1e-5, 0, np.inf),
+            (0, 0, np.nan),
+            (1e-7, 1e-7, 1),
+        ],
+    )
+    def test_smooth(
+        self,
+        value: float,
+        smooth_nr: float,
+        smooth_dr: float,
+        expected: float,
+    ):
+        """
+        Test values in extreme cases where variances are all zero.
+
+        :param value: value for input.
+        :param smooth_nr: constant for numerator.
+        :param smooth_dr: constant for denominator.
+        :param expected: target value.
+        """
+        kernel_size = 5
+        mid = kernel_size // 2
+        shape = (1, kernel_size, kernel_size, kernel_size, 1)
+        y_true = tf.ones(shape=shape) * value
+        y_pred = tf.ones(shape=shape) * value
+
+        got = image.LocalNormalizedCrossCorrelation(
+            kernel_size=kernel_size,
+            smooth_nr=smooth_nr,
+            smooth_dr=smooth_dr,
+        ).calc_ncc(
             y_true,
             y_pred,
         )
+        got = got[0, mid, mid, mid, 0]
+        expected = tf.constant(expected)
+        assert is_equal_tf(got, expected)
+
+    @pytest.mark.parametrize(
+        "kernel_type",
+        ["rectangular", "gaussian", "triangular"],
+    )
+    @pytest.mark.parametrize(
+        "kernel_size",
+        [3, 5, 7],
+    )
+    def test_exact_value(self, kernel_type, kernel_size):
+        """
+        Test the exact value at the center of a cube.
+
+        :param kernel_type: name of kernel.
+        :param kernel_size: size of the kernel and the cube.
+        """
+        # init
+        mid = kernel_size // 2
+        tf.random.set_seed(0)
+        y_true = tf.random.uniform(shape=(1, kernel_size, kernel_size, kernel_size, 1))
+        y_pred = tf.random.uniform(shape=(1, kernel_size, kernel_size, kernel_size, 1))
+        loss = image.LocalNormalizedCrossCorrelation(
+            kernel_type=kernel_type, kernel_size=kernel_size
+        )
+
+        # obtained value
+        got = loss.calc_ncc(y_true=y_true, y_pred=y_pred)
+        got = got[0, mid, mid, mid, 0]  # center voxel
+
+        # target value
+        kernel_3d = (
+            loss.kernel[:, None, None]
+            * loss.kernel[None, :, None]
+            * loss.kernel[None, None, :]
+        )
+        kernel_3d = kernel_3d[None, :, :, :, None]
+
+        y_true_mean = tf.reduce_sum(y_true * kernel_3d) / loss.kernel_vol
+        y_true_normalized = y_true - y_true_mean
+        y_true_var = tf.reduce_sum(y_true_normalized ** 2 * kernel_3d)
+
+        y_pred_mean = tf.reduce_sum(y_pred * kernel_3d) / loss.kernel_vol
+        y_pred_normalized = y_pred - y_pred_mean
+        y_pred_var = tf.reduce_sum(y_pred_normalized ** 2 * kernel_3d)
+
+        cross = tf.reduce_sum(y_true_normalized * y_pred_normalized * kernel_3d)
+        expected = (cross ** 2 + EPS) / (y_pred_var * y_true_var + EPS)
+
+        # check
         assert is_equal_tf(got, expected)
 
-    def test_error(self):
-        y = np.ones(shape=(3, 3, 3, 3))
+    def test_kernel_error(self):
+        """Test the error message when using wrong kernel."""
         with pytest.raises(ValueError) as err_info:
-            image.LocalNormalizedCrossCorrelation(kernel_type="constant").call(y, y)
+            image.LocalNormalizedCrossCorrelation(kernel_type="constant")
         assert "Wrong kernel_type constant for LNCC loss type." in str(err_info.value)
 
     def test_get_config(self):
+        """Test the config is saved correctly."""
         got = image.LocalNormalizedCrossCorrelation().get_config()
         expected = dict(
             kernel_size=9,
             kernel_type="rectangular",
             reduction=tf.keras.losses.Reduction.SUM,
             name="LocalNormalizedCrossCorrelation",
+            smooth_nr=1e-5,
+            smooth_dr=1e-5,
         )
         assert got == expected
 
diff --git a/test/unit/test_loss_label.py b/test/unit/test_loss_label.py
index ba43f2ec..60007d68 100644
--- a/test/unit/test_loss_label.py
+++ b/test/unit/test_loss_label.py
@@ -12,6 +12,7 @@ import pytest
 import tensorflow as tf
 
 import deepreg.loss.label as label
+from deepreg.constant import EPS
 
 
 class TestMultiScaleLoss:
@@ -91,11 +92,11 @@ class TestCrossEntropy:
     @pytest.mark.parametrize(
         "binary,background_weight,scales,expected",
         [
-            (True, 0.0, None, -np.log(1.0e-7)),
-            (False, 0.0, None, -0.6 * np.log(0.3)),
-            (False, 0.2, None, -0.48 * np.log(0.3) - 0.08 * np.log(0.7)),
-            (False, 0.2, [0, 0], -0.48 * np.log(0.3) - 0.08 * np.log(0.7)),
-            (False, 0.2, [0, 1], 0.5239637),
+            (True, 0.0, None, -np.log(EPS)),
+            (False, 0.0, None, -0.6 * np.log(0.3 + EPS)),
+            (False, 0.2, None, -0.48 * np.log(0.3 + EPS) - 0.08 * np.log(0.7 + EPS)),
+            (False, 0.2, [0, 0], -0.48 * np.log(0.3 + EPS) - 0.08 * np.log(0.7 + EPS)),
+            (False, 0.2, [0, 1], 0.5239465),
         ],
     )
     def test_call(self, y_true, y_pred, binary, background_weight, scales, expected):
@@ -135,7 +136,7 @@ class TestJaccardIndex:
             (True, None, 0),
             (False, None, 0.25),
             (False, [0, 0], 0.25),
-            (False, [0, 1], 0.17484076),
+            (False, [0, 1], 0.17485845),
         ],
     )
     def test_call(self, y_true, y_pred, binary, scales, expected):
diff --git a/test/unit/util.py b/test/unit/util.py
index 02866779..d16c5561 100644
--- a/test/unit/util.py
+++ b/test/unit/util.py
@@ -3,12 +3,14 @@ from typing import List, Union
 import numpy as np
 import tensorflow as tf
 
+from deepreg.constant import EPS
+
 
 def is_equal_np(
-    x: Union[np.ndarray, List], y: Union[np.ndarray, List], atol: float = 1.0e-7
+    x: Union[np.ndarray, List], y: Union[np.ndarray, List], atol: float = EPS
 ) -> bool:
     """
-    Check if two numpy arrays are identical.
+    Check if two numpy arrays are identical within a tolerance.
 
     :param x:
     :param y:
@@ -17,16 +19,29 @@ def is_equal_np(
     """
     x = np.asarray(x, dtype=np.float32)
     y = np.asarray(y, dtype=np.float32)
-    return x.shape == y.shape and np.all(np.isclose(x, y, atol=atol))
+
+    # check shape
+    if x.shape != y.shape:
+        return False
+
+    # check nan values
+    # support case some values are nan
+    if np.any(np.isnan(x) != np.isnan(y)):
+        return False
+    x = np.nan_to_num(x)
+    y = np.nan_to_num(y)
+
+    # check values
+    return np.all(np.isclose(x, y, atol=atol))
 
 
 def is_equal_tf(
     x: Union[tf.Tensor, np.ndarray, List],
     y: Union[tf.Tensor, np.ndarray, List],
-    atol: float = 1.0e-7,
+    atol: float = EPS,
 ) -> bool:
     """
-    Check if two tf tensors are identical.
+    Check if two tf tensors are identical within a tolerance.
 
     :param x:
     :param y:
@@ -35,4 +50,4 @@ def is_equal_tf(
     """
     x = tf.cast(x, dtype=tf.float32).numpy()
     y = tf.cast(y, dtype=tf.float32).numpy()
-    return x.shape == y.shape and np.all(np.isclose(x, y, atol=atol))
+    return is_equal_np(x=x, y=y, atol=atol)