Gb/trainable pool

phygnn/layers/custom_layers.py

@@ -136,7 +136,7 @@ class GaussianAveragePooling2D(tf.keras.layers.Layer):
     convolution window that limits the area of effect"""
 
     def __init__(self, pool_size, strides=None, padding='valid', sigma=1,
-                 **kwargs):
+                 trainable=True, **kwargs):
         """
         Parameters
         ----------
@@ -154,47 +154,56 @@ def __init__(self, pool_size, strides=None, padding='valid', sigma=1,
             same height/width dimension as the input.
         sigma : float
             Sigma parameter for gaussian distribution
+        trainable : bool
+            Flag for whether sigma is trainable weight or not.
         kwargs : dict
             Extra kwargs for tf.keras.layers.Layer
         """
 
         super().__init__(**kwargs)
         assert isinstance(pool_size, int), 'pool_size must be int!'
-        self._pool_size = pool_size
-        self._strides = strides
-        self._padding = padding.upper()
-        self._sigma = sigma
-
-        target_shape = (self._pool_size, self._pool_size, 1, 1)
-        self._kernel = self._make_2D_gaussian_kernel(self._pool_size,
-                                                     self._sigma)
-        self._kernel = np.expand_dims(self._kernel, -1)
-        self._kernel = np.expand_dims(self._kernel, -1)
-        assert self._kernel.shape == target_shape
-        self._kernel = tf.convert_to_tensor(self._kernel, dtype=tf.float32)
+        self.pool_size = pool_size
+        self.strides = strides
+        self.padding = padding.upper()
+        self.trainable = trainable
+        self.sigma = sigma
 
-    @staticmethod
-    def _make_2D_gaussian_kernel(edge_len, sigma=1.):
-        """Creates 2D gaussian kernel with side length `edge_len` and a sigma
-        of `sigma`
+    def build(self, input_shape):
+        """Custom implementation of the tf layer build method.
+
+        Initializes the trainable sigma variable
 
         Parameters
         ----------
-        edge_len : int
-            Edge size of the kernel
-        sigma : float
-            Sigma parameter for gaussian distribution
+        input_shape : tuple
+            Shape tuple of the input
+        """
+        if not any(self.weights):
+            init = tf.keras.initializers.Constant(value=self.sigma)
+            self.sigma = self.add_weight("sigma", shape=[1],
+                                         trainable=self.trainable,
+                                         dtype=tf.float32,
+                                         initializer=init)
+
+    def make_kernel(self):
+        """Creates 2D gaussian kernel with side length `self.pool_size` and a
+        sigma of `sigma`
 
         Returns
         -------
         kernel : np.ndarray
-            2D kernel with shape (edge_len, edge_len)
+            2D kernel with shape (self.pool_size, self.pool_size)
         """
-        ax = np.linspace(-(edge_len - 1) / 2., (edge_len - 1) / 2., edge_len)
-        gauss = np.exp(-0.5 * np.square(ax) / np.square(sigma))
-        kernel = np.outer(gauss, gauss)
-        kernel = kernel / np.sum(kernel)
-        return kernel.astype(np.float32)
+        ax = tf.linspace(-(self.pool_size - 1) / 2.,
+                         (self.pool_size - 1) / 2.,
+                         self.pool_size)
+        gauss = tf.math.exp(-0.5 * tf.math.square(ax)
+                            / tf.math.square(self.sigma))
+        kernel = tf.expand_dims(gauss, 0) * tf.expand_dims(gauss, -1)
+        kernel = kernel / tf.math.reduce_sum(kernel)
+        kernel = tf.expand_dims(kernel, -1)
+        kernel = tf.expand_dims(kernel, -1)
+        return kernel
 
     def get_config(self):
         """Implementation of get_config method from tf.keras.layers.Layer for
@@ -206,10 +215,11 @@ def get_config(self):
         """
         config = super().get_config().copy()
         config.update({
-            'pool_size': self._pool_size,
-            'strides': self._strides,
-            'padding': self._padding,
-            'sigma': self._sigma,
+            'pool_size': self.pool_size,
+            'strides': self.strides,
+            'padding': self.padding,
+            'trainable': self.trainable,
+            'sigma': float(self.sigma),
         })
         return config
 
@@ -226,12 +236,15 @@ def call(self, x):
         x : tf.Tensor
             Output tensor operated on by the specified function
         """
+
+        kernel = self.make_kernel()
+
         out = []
         for idf in range(x.shape[-1]):
             fslice = slice(idf, idf + 1)
-            iout = tf.nn.convolution(x[..., fslice], self._kernel,
-                                     strides=self._strides,
-                                     padding=self._padding)
+            iout = tf.nn.convolution(x[..., fslice], kernel,
+                                     strides=self.strides,
+                                     padding=self.padding)
             out.append(iout)
         out = tf.concat(out, -1, name='concat')
         return out

tests/test_layers.py

@@ -467,10 +467,10 @@ def test_gaussian_pooling():
     for stdev in [1, 2]:
         layer = GaussianAveragePooling2D(pool_size=5, strides=1, sigma=stdev)
         _ = layer(np.ones((24, 100, 100, 35)))
-        kernel = layer._kernel.numpy()
+        kernel = layer.make_kernel().numpy()
         kernels.append(kernel)
 
-        assert kernel[:, :, 0, 0].sum() == 1
+        assert np.allclose(kernel[:, :, 0, 0].sum(), 1, rtol=1e-4)
         assert kernel[2, 2, 0, 0] == kernel.max()
         assert kernel[0, 0, 0, 0] == kernel.min()
         assert kernel[-1, -1, 0, 0] == kernel.min()
@@ -485,7 +485,7 @@ def test_gaussian_pooling():
                            normalize=False)
     x_in = np.random.uniform(0, 1, (1, 12, 12, 3))
     out1 = model1.predict(x_in)
-    kernel1 = model1.layers[0]._kernel[:, :, 0, 0].numpy()
+    kernel1 = model1.layers[0].make_kernel()[:, :, 0, 0].numpy()
 
     for idf in range(out1.shape[-1]):
         test = (x_in[0, :, :, idf] * kernel1).sum()
@@ -499,7 +499,7 @@ def test_gaussian_pooling():
         model1.save_model(model_path)
         model2 = TfModel.load(model_path)
 
-        kernel2 = model2.layers[0]._kernel[:, :, 0, 0].numpy()
+        kernel2 = model2.layers[0].make_kernel()[:, :, 0, 0].numpy()
         out2 = model2.predict(x_in)
         assert np.allclose(kernel1, kernel2)
         assert np.allclose(out1, out2)
@@ -509,6 +509,52 @@ def test_gaussian_pooling():
         _ = model2.predict(x_in)
 
 
+def test_gaussian_pooling_train():
+    """Test the trainable sigma functionality of the gaussian average pool"""
+    pool_size = 5
+    xtrain = np.random.uniform(0, 1, (10, pool_size, pool_size, 1))
+    ytrain = np.random.uniform(0, 1, (10, 1, 1, 1))
+    hidden_layers = [{'class': 'GaussianAveragePooling2D',
+                      'pool_size': pool_size, 'trainable': False,
+                      'strides': 1, 'padding': 'valid', 'sigma': 2}]
+
+    model = TfModel.build(['x'], ['y'],
+                          hidden_layers=hidden_layers,
+                          input_layer=False,
+                          output_layer=False,
+                          learning_rate=1e-3,
+                          normalize=(True, True))
+    model.layers[0].build(xtrain.shape)
+    assert len(model.layers[0].trainable_weights) == 0
+
+    hidden_layers[0]['trainable'] = True
+    model = TfModel.build(['x'], ['y'],
+                          hidden_layers=hidden_layers,
+                          input_layer=False,
+                          output_layer=False,
+                          learning_rate=1e-3,
+                          normalize=(True, True))
+    model.layers[0].build(xtrain.shape)
+    assert len(model.layers[0].trainable_weights) == 1
+
+    layer = model.layers[0]
+    sigma1 = float(layer.sigma)
+    kernel1 = layer.make_kernel().numpy().copy()
+    model.train_model(xtrain, ytrain, epochs=10)
+    sigma2 = float(layer.sigma)
+    kernel2 = layer.make_kernel().numpy().copy()
+
+    assert not np.allclose(sigma1, sigma2)
+    assert not np.allclose(kernel1, kernel2)
+
+    with TemporaryDirectory() as td:
+        model_path = os.path.join(td, 'test_model')
+        model.save_model(model_path)
+        model2 = TfModel.load(model_path)
+
+    assert np.allclose(model.predict(xtrain), model2.predict(xtrain))
+
+
 def test_siglin():
     """Test the sigmoid linear layer"""
     n_points = 1000
@@ -532,7 +578,7 @@ def test_logtransform():
 
     lt = LogTransform(adder=1)
     ilt = LogTransform(adder=1, inverse=True)
-    x = np.random.uniform(0, 10, (n_points + 1, 2))
+    x = np.random.uniform(0.01, 10, (n_points + 1, 2))
     y = lt(x).numpy()
     xinv = ilt(y).numpy()
     assert not np.isnan(y).any()
@@ -541,7 +587,7 @@ def test_logtransform():
 
     lt = LogTransform(adder=1, idf=1)
     ilt = LogTransform(adder=1, inverse=True, idf=1)
-    x = np.random.uniform(0, 10, (n_points + 1, 2))
+    x = np.random.uniform(0.01, 10, (n_points + 1, 2))
     y = lt(x).numpy()
     xinv = ilt(y).numpy()
     assert np.allclose(x[:, 0], y[:, 0])