def _dilated_conv_layer(self, output_channels, dilation_rate, apply_relu,
name):
"""Create a dilated convolution layer.
Args:
output_channels: int. Number of output channels for each pixel.
dilation_rate: int. Represents how many pixels each stride offset will
move. A value of 1 indicates a standard convolution.
apply_relu: bool. If True, a ReLU non-linearlity is added.
name: string. Name for layer.
Returns:
a sonnet Module for a dilated convolution.
"""
layer_components = [
conv.Conv2D(output_channels, [3, 3],
initializers=self._initializers,
regularizers=self._regularizers,
rate=dilation_rate,
name="dilated_conv_" + name),
]
if apply_relu:
layer_components.append(
lambda net: tf.nn.relu(net, name="relu_" + name))
return sequential.Sequential(layer_components, name=name)
def _build(self, images):
"""Build dilation module.
Args:
images: Tensor of shape [batch_size, height, width, depth]
and dtype float32. Represents a set of images with an arbitrary depth.
Note that when using the default initializer, depth must equal
num_output_classes.
Returns:
Tensor of shape [batch_size, height, width, num_output_classes] and dtype
float32. Represents, for each image and pixel, logits for per-class
predictions.
Raises:
IncompatibleShapeError: If images is not rank 4.
ValueError: If model_size is not one of 'basic' or 'large'.
"""
num_classes = self._num_output_classes
if len(images.get_shape()) != 4:
raise base.IncompatibleShapeError(
"'images' must have shape [batch_size, height, width, depth].")
if self.WEIGHTS not in self._initializers:
if self._model_size == self.BASIC:
self._initializers[self.WEIGHTS] = identity_kernel_initializer
elif self._model_size == self.LARGE:
self._initializers[
self.WEIGHTS] = noisy_identity_kernel_initializer(
num_classes)
else:
raise ValueError("Unrecognized model_size: %s" %
self._model_size)
if self.BIASES not in self._initializers:
self._initializers[self.BIASES] = tf.zeros_initializer()
if self._model_size == self.BASIC:
self._conv_modules = [
self._dilated_conv_layer(num_classes, 1, True, "conv1"),
self._dilated_conv_layer(num_classes, 1, True, "conv2"),
self._dilated_conv_layer(num_classes, 2, True, "conv3"),
self._dilated_conv_layer(num_classes, 4, True, "conv4"),
self._dilated_conv_layer(num_classes, 8, True, "conv5"),
self._dilated_conv_layer(num_classes, 16, True, "conv6"),
self._dilated_conv_layer(num_classes, 1, True, "conv7"),
self._dilated_conv_layer(num_classes, 1, False, "conv8"),
]
elif self._model_size == self.LARGE:
self._conv_modules = [
self._dilated_conv_layer(2 * num_classes, 1, True, "conv1"),
self._dilated_conv_layer(2 * num_classes, 1, True, "conv2"),
self._dilated_conv_layer(4 * num_classes, 2, True, "conv3"),
self._dilated_conv_layer(8 * num_classes, 4, True, "conv4"),
self._dilated_conv_layer(16 * num_classes, 8, True, "conv5"),
self._dilated_conv_layer(32 * num_classes, 16, True, "conv6"),
self._dilated_conv_layer(32 * num_classes, 1, True, "conv7"),
self._dilated_conv_layer(num_classes, 1, False, "conv8"),
]
else:
raise ValueError("Unrecognized model_size: %s" % self._model_size)
dilation_mod = sequential.Sequential(self._conv_modules,
name="dilation")
return dilation_mod(images)