# -*- coding: utf-8 -*-
from __future__ import absolute_import, print_function
from six.moves import range
from niftynet.layer import layer_util
from niftynet.layer.activation import ActiLayer
from niftynet.layer.base_layer import TrainableLayer
from niftynet.layer.bn import BNLayer
from niftynet.layer.convolution import ConvLayer, ConvolutionalLayer
from niftynet.layer.dilatedcontext import DilatedTensor
from niftynet.layer.elementwise import ElementwiseLayer
from niftynet.network.base_net import BaseNet
[docs]class HighRes3DNet(BaseNet):
"""
implementation of HighRes3DNet:
Li et al., "On the compactness, efficiency, and representation of 3D
convolutional networks: Brain parcellation as a pretext task", IPMI '17
### Building blocks
{ } - Residual connections: see He et al. "Deep residual learning for
image recognition", in CVPR '16
[CONV] - Convolutional layer in form: Activation(Convolution(X))
where X = input tensor or output of previous layer
and Activation is a function which includes:
a) Batch-Norm
b) Activation Function (ReLu, PreLu, Sigmoid, Tanh etc.)
c) Drop-out layer by sampling random variables from a Bernouilli distribution
if p < 1
[CONV*] - Convolutional layer with no activation function
(r)[D-CONV(d)] - Convolutional layer with dilated convolutions with blocks in
pre-activation mode: D-Convolution(Activation(X))
see He et al., "Identity Mappings in Deep Residual Networks", ECCV '16
dilation factor = d
D-CONV(2) : dilated convolution with dilation factor 2
repeat factor = r
e.g.
(2)[D-CONV(d)] : 2 dilated convolutional layers in a row [D-CONV] -> [D-CONV]
{ (2)[D-CONV(d)] } : 2 dilated convolutional layers within residual block
### Diagram
INPUT --> [CONV] --> { (3)[D-CONV(1)] } --> { (3)[D-CONV(2)] } --> { (3)[D-CONV(4)] } -> [CONV*] -> Loss
### Constraints
- Input image size should be divisible by 8
"""
[docs] def __init__(self,
num_classes,
w_initializer=None,
w_regularizer=None,
b_initializer=None,
b_regularizer=None,
acti_func='prelu',
name='HighRes3DNet'):
"""
:param num_classes: int, number of channels of output
:param w_initializer: weight initialisation for network
:param w_regularizer: weight regularisation for network
:param b_initializer: bias initialisation for network
:param b_regularizer: bias regularisation for network
:param acti_func: activation function to use
:param name: layer name
"""
super(HighRes3DNet, self).__init__(
num_classes=num_classes,
w_initializer=w_initializer,
w_regularizer=w_regularizer,
b_initializer=b_initializer,
b_regularizer=b_regularizer,
acti_func=acti_func,
name=name)
self.layers = [
{'name': 'conv_0', 'n_features': 16, 'kernel_size': 3},
{'name': 'res_1', 'n_features': 16, 'kernels': (3, 3), 'repeat': 3},
{'name': 'res_2', 'n_features': 32, 'kernels': (3, 3), 'repeat': 3},
{'name': 'res_3', 'n_features': 64, 'kernels': (3, 3), 'repeat': 3},
{'name': 'conv_1', 'n_features': 80, 'kernel_size': 1},
{'name': 'conv_2', 'n_features': num_classes, 'kernel_size': 1}]
[docs] def layer_op(self, images, is_training=True, layer_id=-1, **unused_kwargs):
"""
:param images: tensor to input to the network. Size has to be divisible by 8
:param is_training: boolean, True if network is in training mode
:param layer_id: int, index of the layer to return as output
:param unused_kwargs:
:return: output of layer indicated by layer_id
"""
assert layer_util.check_spatial_dims(
images, lambda x: x % 8 == 0)
# go through self.layers, create an instance of each layer
# and plugin data
layer_instances = []
### first convolution layer
params = self.layers[0]
first_conv_layer = ConvolutionalLayer(
n_output_chns=params['n_features'],
kernel_size=params['kernel_size'],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name=params['name'])
flow = first_conv_layer(images, is_training)
layer_instances.append((first_conv_layer, flow))
### resblocks, all kernels dilated by 1 (normal convolution)
params = self.layers[1]
with DilatedTensor(flow, dilation_factor=1) as dilated:
for j in range(params['repeat']):
res_block = HighResBlock(
params['n_features'],
params['kernels'],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='%s_%d' % (params['name'], j))
dilated.tensor = res_block(dilated.tensor, is_training)
layer_instances.append((res_block, dilated.tensor))
flow = dilated.tensor
### resblocks, all kernels dilated by 2
params = self.layers[2]
with DilatedTensor(flow, dilation_factor=2) as dilated:
for j in range(params['repeat']):
res_block = HighResBlock(
params['n_features'],
params['kernels'],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='%s_%d' % (params['name'], j))
dilated.tensor = res_block(dilated.tensor, is_training)
layer_instances.append((res_block, dilated.tensor))
flow = dilated.tensor
### resblocks, all kernels dilated by 4
params = self.layers[3]
with DilatedTensor(flow, dilation_factor=4) as dilated:
for j in range(params['repeat']):
res_block = HighResBlock(
params['n_features'],
params['kernels'],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='%s_%d' % (params['name'], j))
dilated.tensor = res_block(dilated.tensor, is_training)
layer_instances.append((res_block, dilated.tensor))
flow = dilated.tensor
### 1x1x1 convolution layer
params = self.layers[4]
fc_layer = ConvolutionalLayer(
n_output_chns=params['n_features'],
kernel_size=params['kernel_size'],
acti_func=self.acti_func,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name=params['name'])
flow = fc_layer(flow, is_training)
layer_instances.append((fc_layer, flow))
### 1x1x1 convolution layer
params = self.layers[5]
fc_layer = ConvolutionalLayer(
n_output_chns=params['n_features'],
kernel_size=params['kernel_size'],
acti_func=None,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name=params['name'])
flow = fc_layer(flow, is_training)
layer_instances.append((fc_layer, flow))
# set training properties
if is_training:
self._print(layer_instances)
return layer_instances[-1][1]
return layer_instances[layer_id][1]
def _print(self, list_of_layers):
for (op, _) in list_of_layers:
print(op)
[docs]class HighResBlock(TrainableLayer):
"""
This class defines a high-resolution block with residual connections
kernels
- specify kernel sizes of each convolutional layer
- e.g.: kernels=(5, 5, 5) indicate three conv layers of kernel_size 5
with_res
- whether to add residual connections to bypass the conv layers
"""
[docs] def __init__(self,
n_output_chns,
kernels=(3, 3),
acti_func='relu',
w_initializer=None,
w_regularizer=None,
with_res=True,
name='HighResBlock'):
"""
:param n_output_chns: int, number of output channels
:param kernels: list of layer kernel sizes
:param acti_func: activation function to use
:param w_initializer: weight initialisation for network
:param w_regularizer: weight regularisation for network
:param with_res: boolean, set to True if residual connection are to use
:param name: layer name
"""
super(HighResBlock, self).__init__(name=name)
self.n_output_chns = n_output_chns
if hasattr(kernels, "__iter__"): # a list of layer kernel_sizes
self.kernels = kernels
else: # is a single number (indicating single layer)
self.kernels = [kernels]
self.acti_func = acti_func
self.with_res = with_res
self.initializers = {'w': w_initializer}
self.regularizers = {'w': w_regularizer}
[docs] def layer_op(self, input_tensor, is_training):
"""
:param input_tensor: tensor, input to the network
:param is_training: boolean, True if network is in training mode
:return: tensor, output of the residual block
"""
output_tensor = input_tensor
for (i, k) in enumerate(self.kernels):
# create parameterised layers
bn_op = BNLayer(regularizer=self.regularizers['w'],
name='bn_{}'.format(i))
acti_op = ActiLayer(func=self.acti_func,
regularizer=self.regularizers['w'],
name='acti_{}'.format(i))
conv_op = ConvLayer(n_output_chns=self.n_output_chns,
kernel_size=k,
stride=1,
w_initializer=self.initializers['w'],
w_regularizer=self.regularizers['w'],
name='conv_{}'.format(i))
# connect layers
output_tensor = bn_op(output_tensor, is_training)
output_tensor = acti_op(output_tensor)
output_tensor = conv_op(output_tensor)
# make residual connections
if self.with_res:
output_tensor = ElementwiseLayer('SUM')(output_tensor, input_tensor)
return output_tensor
