Issue #557: refactoring models

config/test/affine.yaml

@@ -0,0 +1,32 @@
+train:
+  method: "ddf" # ddf / dvf / conditional
+  backbone:
+    name: "global"
+    num_channel_initial: 1
+    extract_levels: [0, 1, 2, 3, 4]
+  loss:
+    image:
+      name: "lncc"
+      weight: 0.1
+    label:
+      weight: 1.0
+      name: "dice"
+      scales: [0, 1, 2, 4, 8, 16, 32]
+    regularization:
+      weight: 0.5
+      name: "bending"
+  preprocess:
+    batch_size: 2
+    shuffle_buffer_num_batch: 1
+  optimizer:
+    name: "adam"
+    adam:
+      learning_rate: 1.0e-5
+    sgd:
+      learning_rate: 1.0e-4
+      momentum: 0.9
+    rms:
+      learning_rate: 1.0e-4
+      momentum: 0.9
+  epochs: 2
+  save_period: 2

config/unpaired_labeled_ddf.yaml

@@ -10,7 +10,7 @@ dataset:
 
 train:
   # define neural network structure
-  method: "ddf" # options include "ddf", "dvf", "conditional" and "affine"
+  method: "ddf" # options include "ddf", "dvf", "conditional"
   backbone:
     name: "local" # options include "local", "unet" and "global" - use "global" when method=="affine"
     num_channel_initial: 1 # number of initial channel in local net, controls the size of the network

deepreg/dataset/loader/grouped_loader.py

@@ -5,6 +5,7 @@
 Read https://deepreg.readthedocs.io/en/latest/api/loader.html#module-deepreg.dataset.loader.grouped_loader for more details.
 """
 import random
+from copy import deepcopy
 from typing import List
 
 from deepreg.dataset.loader.interface import (
@@ -264,7 +265,7 @@ def sample_index_generator(self):
         else:
             # sample indices are pre-calculated
             assert self.sample_indices is not None
-            sample_indices = self.sample_indices.copy()
+            sample_indices = deepcopy(self.sample_indices)
             rnd.shuffle(sample_indices)  # shuffle in place
             for sample_index in sample_indices:
                 group_index1, image_index1, group_index2, image_index2 = sample_index

deepreg/model/__init__.py

@@ -1,3 +1,4 @@
 # flake8: noqa
 import deepreg.model.backbone
 import deepreg.model.loss
+import deepreg.model.network

deepreg/model/backbone/global_net.py

@@ -70,7 +70,9 @@ def __init__(
             units=12, bias_initializer=self.transform_initial
         )
 
-    def call(self, inputs: tf.Tensor, training=None, mask=None) -> tf.Tensor:
+    def call(
+        self, inputs: tf.Tensor, training=None, mask=None
+    ) -> (tf.Tensor, tf.Tensor):
         """
         Build GlobalNet graph based on built layers.
 
@@ -88,10 +90,10 @@ def call(self, inputs: tf.Tensor, training=None, mask=None) -> tf.Tensor:
         )  # level E of encoding
 
         # predict affine parameters theta of shape = [batch, 4, 3]
-        self.theta = self._dense_layer(h_out)
-        self.theta = tf.reshape(self.theta, shape=(-1, 4, 3))
+        theta = self._dense_layer(h_out)
+        theta = tf.reshape(theta, shape=(-1, 4, 3))
 
         # warp the reference grid with affine parameters to output a ddf
-        grid_warped = layer_util.warp_grid(self.reference_grid, self.theta)
-        output = grid_warped - self.reference_grid
-        return output
+        grid_warped = layer_util.warp_grid(self.reference_grid, theta)
+        ddf = grid_warped - self.reference_grid
+        return ddf, theta

deepreg/model/loss/image.py

@@ -18,13 +18,13 @@ class SumSquaredDifference(tf.keras.losses.Loss):
 
     def __init__(
         self,
-        reduction: str = tf.keras.losses.Reduction.AUTO,
+        reduction: str = tf.keras.losses.Reduction.SUM,
         name: str = "SumSquaredDifference",
     ):
         """
         Init.
 
-        :param reduction: using AUTO reduction,
+        :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
         """
@@ -57,15 +57,15 @@ def __init__(
         self,
         num_bins: int = 23,
         sigma_ratio: float = 0.5,
-        reduction: str = tf.keras.losses.Reduction.AUTO,
+        reduction: str = tf.keras.losses.Reduction.SUM,
         name: str = "GlobalMutualInformation",
     ):
         """
         Init.
 
         :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 AUTO reduction,
+        :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
         """
@@ -248,15 +248,15 @@ def __init__(
         self,
         kernel_size: int = 9,
         kernel_type: str = "rectangular",
-        reduction: str = tf.keras.losses.Reduction.AUTO,
+        reduction: str = tf.keras.losses.Reduction.SUM,
         name: str = "LocalNormalizedCrossCorrelation",
     ):
         """
         Init.
 
         :param kernel_size: int. Kernel size or kernel sigma for kernel_type='gauss'.
         :param kernel_type: str, rectangular, triangular or gaussian
-        :param reduction: using AUTO reduction,
+        :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
         """

deepreg/model/loss/label.py

@@ -53,15 +53,15 @@ def __init__(
         self,
         scales: Optional[List] = None,
         kernel: str = "gaussian",
-        reduction: str = tf.keras.losses.Reduction.AUTO,
+        reduction: str = tf.keras.losses.Reduction.SUM,
         name: str = "MultiScaleLoss",
     ):
         """
         Init.
 
         :param scales: list of scalars or None, if None, do not apply any scaling.
         :param kernel: gaussian or cauchy.
-        :param reduction: using AUTO reduction,
+        :param reduction: using SUM reduction over batch axis,
             calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
         :param name: str, name of the loss.
         """
@@ -144,7 +144,7 @@ def __init__(
         neg_weight: float = 0.0,
         scales: Optional[List] = None,
         kernel: str = "gaussian",
-        reduction: str = tf.keras.losses.Reduction.AUTO,
+        reduction: str = tf.keras.losses.Reduction.SUM,
         name: str = "DiceScore",
     ):
         """
@@ -154,7 +154,7 @@ def __init__(
         :param neg_weight: weight for negative class.
         :param scales: list of scalars or None, if None, do not apply any scaling.
         :param kernel: gaussian or cauchy.
-        :param reduction: using AUTO reduction,
+        :param reduction: using SUM reduction over batch axis,
             calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
         :param name: str, name of the loss.
         """
@@ -214,7 +214,7 @@ def __init__(
         neg_weight: float = 0.0,
         scales: Optional[List] = None,
         kernel: str = "gaussian",
-        reduction: str = tf.keras.losses.Reduction.AUTO,
+        reduction: str = tf.keras.losses.Reduction.SUM,
         name: str = "CrossEntropy",
     ):
         """
@@ -224,7 +224,7 @@ def __init__(
         :param neg_weight: weight for negative class
         :param scales: list of scalars or None, if None, do not apply any scaling.
         :param kernel: gaussian or cauchy.
-        :param reduction: using AUTO reduction,
+        :param reduction: using SUM reduction over batch axis,
             calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
         :param name: str, name of the loss.
         """
@@ -276,7 +276,7 @@ def __init__(
         binary: bool = False,
         scales: Optional[List] = None,
         kernel: str = "gaussian",
-        reduction: str = tf.keras.losses.Reduction.AUTO,
+        reduction: str = tf.keras.losses.Reduction.SUM,
         name: str = "JaccardIndex",
     ):
         """
@@ -285,7 +285,7 @@ def __init__(
         :param binary: if True, project y_true, y_pred to 0 or 1.
         :param scales: list of scalars or None, if None, do not apply any scaling.
         :param kernel: gaussian or cauchy.
-        :param reduction: using AUTO reduction,
+        :param reduction: using SUM reduction over batch axis,
             calling the loss like `loss(y_true, y_pred)` will return a scalar tensor.
         :param name: str, name of the loss.
         """