Merge pull request #217 from NifTK/se_block

ChannelSELayer, SpatialSELayer, and ChannelSpatialSELayer.

Author: wyli

niftynet/layer/squeeze_excitation.py

@@ -0,0 +1,128 @@
+# -*- coding: utf-8 -*-
+from __future__ import absolute_import, print_function
+
+import tensorflow as tf
+
+from niftynet.layer.base_layer import Layer
+from niftynet.layer.fully_connected import FullyConnectedLayer
+from niftynet.layer.convolution import ConvolutionalLayer
+from niftynet.utilities.util_common import look_up_operations
+
+SUPPORTED_OP = set(['AVG', 'MAX'])
+
+
+class ChannelSELayer(Layer):
+    """
+    Re-implementation of Squeeze-and-Excitation (SE) block described in::
+
+         Hu et al., Squeeze-and-Excitation Networks, arXiv:1709.01507
+    """
+    def __init__(self,
+                 func='AVG',
+                 reduction_ratio=16,
+                 name='channel_squeeze_excitation'):
+        self.func = func.upper()
+        self.reduction_ratio = reduction_ratio
+        super(ChannelSELayer, self).__init__(name=name)
+
+        look_up_operations(self.func, SUPPORTED_OP)
+
+    def layer_op(self, input_tensor):
+        # spatial squeeze
+        input_rank = len(input_tensor.shape)
+        reduce_indices = list(range(input_rank))[1:-1]
+        if self.func == 'AVG':
+            squeeze_tensor = tf.reduce_mean(input_tensor, axis=reduce_indices)
+        elif self.func == 'MAX':
+            squeeze_tensor = tf.reduce_max(input_tensor, axis=reduce_indices)
+        else:
+            raise NotImplementedError("pooling function not supported")
+
+        # channel excitation
+        num_channels = int(squeeze_tensor.shape[-1])
+        reduction_ratio = self.reduction_ratio
+        if num_channels % reduction_ratio != 0:
+            raise ValueError(
+                "reduction ratio incompatible with "
+                "number of input tensor channels")
+
+        num_channels_reduced = num_channels / reduction_ratio
+        fc1 = FullyConnectedLayer(num_channels_reduced,
+                                  with_bias=False,
+                                  with_bn=False,
+                                  acti_func='relu',
+                                  name='se_fc_1')
+        fc2 = FullyConnectedLayer(num_channels,
+                                  with_bias=False,
+                                  with_bn=False,
+                                  acti_func='sigmoid',
+                                  name='se_fc_2')
+
+        fc_out_1 = fc1(squeeze_tensor)
+        fc_out_2 = fc2(fc_out_1)
+
+        while len(fc_out_2.shape) < input_rank:
+            fc_out_2 = tf.expand_dims(fc_out_2, axis=1)
+
+        output_tensor = tf.multiply(input_tensor, fc_out_2)
+
+        return output_tensor
+
+
+class SpatialSELayer(Layer):
+    """
+    Re-implementation of SE block -- squeezing spatially
+    and exciting channel-wise described in::
+
+        Roy et al., Concurrent Spatial and Channel Squeeze & Excitation
+        in Fully Convolutional Networks, arXiv:1803.02579
+
+    """
+    def __init__(self,
+                 name='spatial_squeeze_excitation'):
+        super(SpatialSELayer, self).__init__(name=name)
+
+    def layer_op(self, input_tensor):
+        # channel squeeze
+        conv = ConvolutionalLayer(n_output_chns=1,
+                                  kernel_size=1,
+                                  with_bn=False,
+                                  acti_func='sigmoid',
+                                  name="se_conv")
+
+        squeeze_tensor = conv(input_tensor)
+
+        # spatial excitation
+        output_tensor = tf.multiply(input_tensor, squeeze_tensor)
+
+        return output_tensor
+
+
+class ChannelSpatialSELayer(Layer):
+    """
+    Re-implementation of concurrent spatial and channel
+    squeeze & excitation::
+
+        Roy et al., Concurrent Spatial and Channel Squeeze & Excitation
+        in Fully Convolutional Networks, arXiv:1803.02579
+
+    """
+    def __init__(self,
+                 func='AVG',
+                 reduction_ratio=16,
+                 name='channel_spatial_squeeze_excitation'):
+        self.func = func.upper()
+        self.reduction_ratio = reduction_ratio
+        super(ChannelSpatialSELayer, self).__init__(name=name)
+
+        look_up_operations(self.func, SUPPORTED_OP)
+
+    def layer_op(self, input_tensor):
+        cSE = ChannelSELayer(func=self.func,
+                             reduction_ratio=self.reduction_ratio,
+                             name='cSE')
+        sSE = SpatialSELayer(name='sSE')
+
+        output_tensor = tf.add(cSE(input_tensor), sSE(input_tensor))
+
+        return output_tensor

niftynet/layer/squeeze_excitation_layer.py

@@ -10,7 +10,7 @@
 from niftynet.layer.fully_connected import FCLayer
 from niftynet.layer.base_layer import TrainableLayer
 from niftynet.layer.convolution import ConvolutionalLayer
-from niftynet.layer.squeeze_excitation_layer import SELayer
+from niftynet.layer.squeeze_excitation import ChannelSELayer
 from niftynet.network.base_net import BaseNet
 
 SE_ResNetDesc = namedtuple('SE_ResNetDesc', ['bn', 'fc', 'conv1', 'blocks'])
@@ -98,7 +98,7 @@ def create(self, input_chns):
 
     def layer_op(self, images, is_training=True):
         layers = self.create(images.shape[-1])
-        se=SELayer()
+        se=ChannelSELayer()
         if self.n_output_chns == images.shape[-1]:
             out=layers.conv[0](images, is_training)
             out=layers.conv[1](out, is_training)