Source code for niftynet.layer.residual_unit

# -*- coding: utf-8 -*-
from __future__ import absolute_import, print_function

from niftynet.layer.activation import ActiLayer as Acti
from niftynet.layer.base_layer import TrainableLayer
from import BNLayer
from niftynet.layer.convolution import ConvolutionalLayer as Conv
from niftynet.layer.elementwise import ElementwiseLayer
from niftynet.utilities.util_common import look_up_operations

SUPPORTED_OP = set(['original', 'conv_bn_acti', 'acti_conv_bn', 'bn_acti_conv'])

[docs]class ResidualUnit(TrainableLayer):
[docs] def __init__(self, n_output_chns=1, kernel_size=3, dilation=1, acti_func='relu', w_initializer=None, w_regularizer=None, moving_decay=0.9, eps=1e-5, type_string='bn_acti_conv', name='res-downsample'): """ Implementation of residual unit presented in: [1] He et al., Identity mapping in deep residual networks, ECCV 2016 [2] He et al., Deep residual learning for image recognition, CVPR 2016 The possible types of connections are:: 'original': residual unit presented in [2] 'conv_bn_acti': ReLU before addition presented in [1] 'acti_conv_bn': ReLU-only pre-activation presented in [1] 'bn_acti_conv': full pre-activation presented in [1] [1] recommends 'bn_acti_conv' :param n_output_chns: number of output feature channels if this doesn't match the input, a 1x1 projection will be created. :param kernel_size: :param dilation: :param acti_func: :param w_initializer: :param w_regularizer: :param moving_decay: :param eps: :param type_string: :param name: """ super(TrainableLayer, self).__init__(name=name) self.type_string = look_up_operations(type_string.lower(), SUPPORTED_OP) self.acti_func = acti_func self.conv_param = {'w_initializer': w_initializer, 'w_regularizer': w_regularizer, 'kernel_size': kernel_size, 'dilation': dilation, 'n_output_chns': n_output_chns} self.bn_param = {'regularizer': w_regularizer, 'moving_decay': moving_decay, 'eps': eps}
[docs] def layer_op(self, inputs, is_training=True): """ The general connections is:: (inputs)--o-conv_0--conv_1-+-- (outputs) | | o----------------o ``conv_0``, ``conv_1`` layers are specified by ``type_string``. """ conv_flow = inputs # batch normalisation layers bn_0 = BNLayer(**self.bn_param) bn_1 = BNLayer(**self.bn_param) # activation functions //regularisers? acti_0 = Acti(func=self.acti_func) acti_1 = Acti(func=self.acti_func) # convolutions conv_0 = Conv(acti_func=None, with_bias=False, feature_normalization=None, **self.conv_param) conv_1 = Conv(acti_func=None, with_bias=False, feature_normalization=None, **self.conv_param) if self.type_string == 'original': conv_flow = acti_0(bn_0(conv_0(conv_flow), is_training)) conv_flow = bn_1(conv_1(conv_flow), is_training) conv_flow = ElementwiseLayer('SUM')(conv_flow, inputs) conv_flow = acti_1(conv_flow) return conv_flow if self.type_string == 'conv_bn_acti': conv_flow = acti_0(bn_0(conv_0(conv_flow), is_training)) conv_flow = acti_1(bn_1(conv_1(conv_flow), is_training)) return ElementwiseLayer('SUM')(conv_flow, inputs) if self.type_string == 'acti_conv_bn': conv_flow = bn_0(conv_0(acti_0(conv_flow)), is_training) conv_flow = bn_1(conv_1(acti_1(conv_flow)), is_training) return ElementwiseLayer('SUM')(conv_flow, inputs) if self.type_string == 'bn_acti_conv': conv_flow = conv_0(acti_0(bn_0(conv_flow, is_training))) conv_flow = conv_1(acti_1(bn_1(conv_flow, is_training))) return ElementwiseLayer('SUM')(conv_flow, inputs) raise ValueError('Unknown type string')