def _build(self):
"""Builds block according to the arguments."""
filters = self._block_args.input_filters * self._block_args.expand_ratio
if self._block_args.expand_ratio != 1:
# Expansion phase:
self._expand_conv = tf.layers.Conv2D(
filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=conv_kernel_initializer,
padding='same',
data_format=self._data_format,
use_bias=False)
self._bn0 = batchnorm(axis=self._channel_axis,
momentum=self._batch_norm_momentum,
epsilon=self._batch_norm_epsilon)
kernel_size = self._block_args.kernel_size
# Depth-wise convolution phase:
self._depthwise_conv = utils.DepthwiseConv2D(
[kernel_size, kernel_size],
strides=self._block_args.strides,
depthwise_initializer=conv_kernel_initializer,
padding='same',
data_format=self._data_format,
use_bias=False)
self._bn1 = batchnorm(axis=self._channel_axis,
momentum=self._batch_norm_momentum,
epsilon=self._batch_norm_epsilon)
if self._has_se:
num_reduced_filters = max(
1,
int(self._block_args.input_filters *
self._block_args.se_ratio))
# Squeeze and Excitation layer.
self._se_reduce = tf.layers.Conv2D(
num_reduced_filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=conv_kernel_initializer,
padding='same',
data_format=self._data_format,
use_bias=True)
self._se_expand = tf.layers.Conv2D(
filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=conv_kernel_initializer,
padding='same',
data_format=self._data_format,
use_bias=True)
# Output phase:
filters = self._block_args.output_filters
self._project_conv = tf.layers.Conv2D(
filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=conv_kernel_initializer,
padding='same',
data_format=self._data_format,
use_bias=False)
self._bn2 = batchnorm(axis=self._channel_axis,
momentum=self._batch_norm_momentum,
epsilon=self._batch_norm_epsilon)
def _build(self):
"""Builds block according to the arguments."""
if self._block_args.super_pixel == 1:
self._superpixel = tf.layers.Conv2D(
self._block_args.input_filters,
kernel_size=[2, 2],
strides=[2, 2],
kernel_initializer=conv_kernel_initializer,
padding='same',
data_format=self._data_format,
use_bias=False)
self._bnsp = self._batch_norm(axis=self._channel_axis,
momentum=self._batch_norm_momentum,
epsilon=self._batch_norm_epsilon)
filters = self._block_args.input_filters * self._block_args.expand_ratio
kernel_size = self._block_args.kernel_size
if self._block_args.fused_conv:
# If use fused mbconv, skip expansion and use regular conv.
self._fused_conv = tf.layers.Conv2D(
filters, [kernel_size, kernel_size],
strides=self._block_args.strides,
kernel_initializer=conv_kernel_initializer,
padding='same',
data_format=self._data_format,
use_bias=False)
else:
# Otherwise, first apply expansion and then apply depthwise conv.
if self._block_args.expand_ratio != 1:
# Expansion phase.
self._expand_conv = tf.layers.Conv2D(
filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=conv_kernel_initializer,
padding='same',
data_format=self._data_format,
use_bias=False)
self._bn0 = self._batch_norm(
axis=self._channel_axis,
momentum=self._batch_norm_momentum,
epsilon=self._batch_norm_epsilon)
# Depth-wise convolution phase:
self._depthwise_conv = utils.DepthwiseConv2D(
[kernel_size, kernel_size],
strides=self._block_args.strides,
depthwise_initializer=conv_kernel_initializer,
padding='same',
data_format=self._data_format,
use_bias=False)
self._bn1 = self._batch_norm(axis=self._channel_axis,
momentum=self._batch_norm_momentum,
epsilon=self._batch_norm_epsilon)
if self._has_se:
num_reduced_filters = max(
1,
int(self._block_args.input_filters *
self._block_args.se_ratio))
# Squeeze and Excitation layer.
self._se_reduce = tf.layers.Conv2D(
num_reduced_filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=conv_kernel_initializer,
padding='same',
data_format=self._data_format,
use_bias=True)
self._se_expand = tf.layers.Conv2D(
filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=conv_kernel_initializer,
padding='same',
data_format=self._data_format,
use_bias=True)
# Output phase:
filters = self._block_args.output_filters
self._project_conv = tf.layers.Conv2D(
filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=conv_kernel_initializer,
padding='same',
data_format=self._data_format,
use_bias=False)
self._bn2 = self._batch_norm(axis=self._channel_axis,
momentum=self._batch_norm_momentum,
epsilon=self._batch_norm_epsilon)
def build(self, input_shape):
assert len(input_shape) == 2
assert len(input_shape[0]) == 4 and len(input_shape[1]) == 4
# Assertion for high inputs
assert input_shape[0][1] // 2 >= self.kernel_size[0]
assert input_shape[0][2] // 2 >= self.kernel_size[1]
# Assertion for low inputs
assert input_shape[0][1] // input_shape[1][1] == 2
assert input_shape[0][2] // input_shape[1][2] == 2
# channels last for TensorFlow
assert K.image_data_format() == "channels_last"
# input channels
high_in = int(input_shape[0][3])
low_in = int(input_shape[1][3])
if self.strided:
self._pool_strd1 = layers.AveragePooling2D(2)
self._pool_strd2 = layers.AveragePooling2D(2)
self._pool_strd3 = layers.AveragePooling2D(2)
self._pool_strd4 = layers.AveragePooling2D(2)
self._upsample_strd3 = self.pass_layer
else:
self._pool_strd1 = self.pass_layer
self._pool_strd2 = self.pass_layer
self._pool_strd3 = self.pass_layer
self._upsample_strd3 = upsample_tpu
self._pool_strd4 = self.pass_layer
# High -> High conv
if self.use_depthwise:
self.high_to_high = utils.DepthwiseConv2D(self.kernel_size,
strides=(1, 1),
depthwise_initializer=self.kernel_initializer,
padding=self.padding,
use_bias=self.use_bias)
# Low -> Low conv
self.low_to_low = utils.DepthwiseConv2D(self.kernel_size,
strides=(1, 1),
depthwise_initializer=self.kernel_initializer,
padding=self.padding,
use_bias=self.use_bias)
self.high_to_low = self.return_none
self.low_to_high = self.return_none
self.add_highs = self.return_first
self.add_lows = self.return_first
self._upsample_strd3 = self.return_none
else:
self.high_to_high = tf.layers.Conv2D(self.high_channels,
kernel_size=self.kernel_size,
strides=(1, 1),
kernel_initializer=self.kernel_initializer,
padding=self.padding,
use_bias=self.use_bias)
# Low -> Low conv
self.low_to_low = tf.layers.Conv2D(self.low_channels,
kernel_size=self.kernel_size,
strides=(1, 1),
kernel_initializer=self.kernel_initializer,
padding=self.padding,
use_bias=False)
# High -> Low conv
self.high_to_low = tf.layers.Conv2D(self.low_channels,
kernel_size=self.kernel_size,
strides=(1, 1),
kernel_initializer=self.kernel_initializer,
padding=self.padding,
use_bias=self.use_bias)
# Low -> High conv
self.low_to_high = tf.layers.Conv2D(self.high_channels,
kernel_size=self.kernel_size,
strides=(1, 1),
kernel_initializer=self.kernel_initializer,
padding=self.padding,
use_bias=False)
self.add_highs = self.add_values
self.add_lows = self.add_values
super().build(input_shape)
def build(self, input_shape):
assert len(input_shape) == 2
assert len(input_shape[0]) == 4 and len(input_shape[1]) == 4
# Assertion for high inputs
assert input_shape[0][1] // 2 >= self.kernel_size[0]
assert input_shape[0][2] // 2 >= self.kernel_size[1]
# Assertion for low inputs
assert input_shape[0][1] // input_shape[1][1] == 2
assert input_shape[0][2] // input_shape[1][2] == 2
# channels last for TensorFlow
assert K.image_data_format() == "channels_last"
# input channels
high_in = int(input_shape[0][3])
low_in = int(input_shape[1][3])
self.avg_pooling = tf.keras.layers.GlobalAveragePooling2D(
data_format="channels_last")
# High -> High conv
if self.use_depthwise:
self.high_to_high = utils.DepthwiseConv2D(
self.kernel_size,
strides=(1, 1),
depthwise_initializer=self.kernel_initializer,
padding=self.padding,
use_bias=self.use_bias)
# Low -> Low conv
self.low_to_low = utils.DepthwiseConv2D(
self.kernel_size,
strides=(1, 1),
depthwise_initializer=self.kernel_initializer,
padding=self.padding,
use_bias=self.use_bias)
self.high_to_low = None
self.low_to_high = None
else:
self.high_to_high = tf.layers.Conv2D(
self.high_channels,
kernel_size=self.kernel_size,
strides=(1, 1),
kernel_initializer=self.kernel_initializer,
padding=self.padding,
use_bias=self.use_bias)
# Low -> Low conv
self.low_to_low = tf.layers.Conv2D(
self.low_channels,
kernel_size=self.kernel_size,
strides=(1, 1),
kernel_initializer=self.kernel_initializer,
padding=self.padding,
use_bias=False)
# High -> Low conv
self.high_to_low = tf.layers.Conv2D(
self.low_channels,
kernel_size=self.kernel_size,
strides=(1, 1),
kernel_initializer=self.kernel_initializer,
padding=self.padding,
use_bias=self.use_bias)
# Low -> High conv
self.low_to_high = tf.layers.Conv2D(
self.high_channels,
kernel_size=self.kernel_size,
strides=(1, 1),
kernel_initializer=self.kernel_initializer,
padding=self.padding,
use_bias=False)
super().build(input_shape)
def _build(self):
"""Builds MnasNet block according to the arguments."""
filters = self._block_args.input_filters * self._block_args.expand_ratio
if self._block_args.expand_ratio != 1:
# Expansion phase:
self._expand_conv = _get_conv2d(
filters=filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=conv_kernel_initializer,
padding='same',
use_bias=False,
data_format=self._data_format,
use_keras=self._use_keras)
# TODO(hongkuny): b/120622234 need to manage update ops directly.
self._bn0 = tf.layers.BatchNormalization(
axis=self._channel_axis,
momentum=self._batch_norm_momentum,
epsilon=self._batch_norm_epsilon,
fused=True)
kernel_size = self._block_args.kernel_size
# Depth-wise convolution phase:
if self._use_keras:
self._depthwise_conv = tf.keras.layers.DepthwiseConv2D(
[kernel_size, kernel_size],
strides=self._block_args.strides,
depthwise_initializer=conv_kernel_initializer,
padding='same',
data_format=self._data_format,
use_bias=False)
else:
self._depthwise_conv = utils.DepthwiseConv2D(
[kernel_size, kernel_size],
strides=self._block_args.strides,
depthwise_initializer=conv_kernel_initializer,
padding='same',
data_format=self._data_format,
use_bias=False)
self._bn1 = tf.layers.BatchNormalization(
axis=self._channel_axis,
momentum=self._batch_norm_momentum,
epsilon=self._batch_norm_epsilon,
fused=True)
if self.has_se:
num_reduced_filters = max(
1,
int(self._block_args.input_filters *
self._block_args.se_ratio))
# Squeeze and Excitation layer.
self._se_reduce = _get_conv2d(
num_reduced_filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=conv_kernel_initializer,
padding='same',
use_bias=True,
data_format=self._data_format,
use_keras=self._use_keras)
self._se_expand = _get_conv2d(
filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=conv_kernel_initializer,
padding='same',
use_bias=True,
data_format=self._data_format,
use_keras=self._use_keras)
# Output phase:
filters = self._block_args.output_filters
self._project_conv = _get_conv2d(
filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=conv_kernel_initializer,
padding='same',
use_bias=False,
data_format=self._data_format,
use_keras=self._use_keras)
self._bn2 = tf.layers.BatchNormalization(
axis=self._channel_axis,
momentum=self._batch_norm_momentum,
epsilon=self._batch_norm_epsilon,
fused=True)
