def depthwise_conv(inputs,
hidden_size,
kernel_size=1,
bias=True,
activation=None,
scope='depthwise_conv',
reuse=None):
with tf.variable_scope(scope, reuse=reuse):
shapes = inputs.shape.as_list()
depthwise_filter = tf.get_variable('depthwise_filter',
(kernel_size, 1, shapes[-1], 1),
dtype=tf.float32,
initializer=he_init())
pointwise_filter = tf.get_variable('pointwise_filter',
(1, 1, shapes[-1], hidden_size),
dtype=tf.float32,
initializer=he_init())
outputs = tf.nn.separable_conv2d(inputs,
depthwise_filter,
pointwise_filter,
strides=(1, 1, 1, 1),
padding='SAME')
if bias:
b = tf.get_variable('bias',
outputs.shape[-1],
initializer=bias_init())
outputs += b
if activation is not None:
return activation(outputs)
else:
return outputs
def conv(x,
channels,
kernel=3,
stride=2,
pad=0,
activation_fn='leaky',
scope='conv_0'):
with tf.variable_scope(scope):
x = tf.pad(x, [[0, 0], [pad, pad], [pad, pad], [0, 0]])
if activation_fn == 'relu':
x = tf.layers.conv2d(
inputs=x,
filters=channels,
kernel_size=kernel,
kernel_initializer=he_init(),
strides=stride,
kernel_regularizer=tf_contrib.layers.l2_regularizer(
scale=0.0001))
else:
x = tf.layers.conv2d(
inputs=x,
filters=channels,
kernel_size=kernel,
strides=stride,
kernel_regularizer=tf_contrib.layers.l2_regularizer(
scale=0.0001))
x = activation(x, activation_fn)
return x
def conv(inputs,
hidden_size,
kernel_size,
bias=True,
activation=tf.nn.relu,
padding='VALID',
scope='conv',
reuse=None):
with tf.variable_scope(scope, reuse=reuse):
in_channels = inputs.get_shape()[-1]
filter_ = tf.get_variable(
'filter',
shape=[1, kernel_size, in_channels, hidden_size],
initializer=he_init(),
dtype=tf.float32)
strides = [1, 1, 1, 1]
out = tf.nn.conv2d(inputs, filter_, strides,
padding) # [N*M, JX, W/filter_stride, d]
if bias:
b = tf.get_variable('bias',
shape=[hidden_size],
initializer=bias_init(),
dtype=tf.float32)
out += b
if activation is not None:
out = activation(out)
out = tf.reduce_max(out, 2)
return out
def deconv(x,
channels,
kernel=3,
stride=2,
activation_fn='leaky',
scope='deconv_0'):
with tf.variable_scope(scope):
if activation_fn == 'relu':
x = tf.layers.conv2d_transpose(
inputs=x,
filters=channels,
kernel_size=kernel,
kernel_initializer=he_init(),
strides=stride,
padding='SAME',
kernel_regularizer=tf_contrib.layers.l2_regularizer(
scale=0.0001))
else:
x = tf.layers.conv2d_transpose(
inputs=x,
filters=channels,
kernel_size=kernel,
strides=stride,
padding='SAME',
kernel_regularizer=tf_contrib.layers.l2_regularizer(
scale=0.0001))
x = activation(x, activation_fn)
return x
def conv(x,
channels,
kernel=3,
stride=2,
pad=0,
normal_weight_init=False,
activation_fn='leaky',
is_training=True,
norm_fn=None,
scope='conv_0'):
with tf.variable_scope(scope):
x = tf.pad(x, [[0, 0], [pad, pad], [pad, pad], [0, 0]], mode=PAD_MODE)
if normal_weight_init:
x = tf.layers.conv2d(
inputs=x,
filters=channels,
kernel_size=kernel,
kernel_initializer=tf.truncated_normal_initializer(
stddev=0.02),
strides=stride,
kernel_regularizer=tf_contrib.layers.l2_regularizer(
scale=0.0001))
else:
if activation_fn == 'relu':
x = tf.layers.conv2d(
inputs=x,
filters=channels,
kernel_size=kernel,
kernel_initializer=he_init(),
strides=stride,
kernel_regularizer=tf_contrib.layers.l2_regularizer(
scale=0.0001))
else:
x = tf.layers.conv2d(
inputs=x,
filters=channels,
kernel_size=kernel,
strides=stride,
kernel_regularizer=tf_contrib.layers.l2_regularizer(
scale=0.0001))
if norm_fn == 'instance':
x = instance_norm(x, 'ins_norm')
if norm_fn == 'batch':
x = batch_norm(x, is_training, 'batch_norm')
x = activation(x, activation_fn)
return x
def resblock(x_init,
channels,
kernel=3,
stride=1,
pad=1,
dropout_ratio=0.0,
is_training=True,
norm_fn='instance',
scope='resblock_0'):
assert norm_fn in ['instance', 'batch', 'weight', 'spectral', None]
with tf.variable_scope(scope):
with tf.variable_scope('res1'):
x = tf.pad(x_init, [[0, 0], [pad, pad], [pad, pad], [0, 0]])
x = tf.layers.conv2d(
inputs=x,
filters=channels,
kernel_size=kernel,
kernel_initializer=he_init(),
strides=stride,
kernel_regularizer=tf_contrib.layers.l2_regularizer(
scale=0.0001))
if norm_fn == 'instance':
x = instance_norm(x, 'res1_instance')
if norm_fn == 'batch':
x = batch_norm(x, is_training, 'res1_batch')
x = relu(x)
with tf.variable_scope('res2'):
x = tf.pad(x, [[0, 0], [pad, pad], [pad, pad], [0, 0]])
x = tf.layers.conv2d(
inputs=x,
filters=channels,
kernel_size=kernel,
strides=stride,
kernel_regularizer=tf_contrib.layers.l2_regularizer(
scale=0.0001))
if norm_fn == 'instance':
x = instance_norm(x, 'res2_instance')
if norm_fn == 'batch':
x = batch_norm(x, is_training, 'res2_batch')
if dropout_ratio > 0.0:
x = tf.layers.dropout(x, rate=dropout_ratio, training=is_training)
return x + x_init
def highway(x,
size=None,
activation=None,
num_layers=2,
dropout=0.0,
scope='highway',
reuse=None):
with tf.variable_scope(scope, reuse=reuse):
if size is None:
size = x.shape.as_list()[-1]
else:
x = fc(x,
size,
weights_initializer=xa_init(),
biases_initializer=bias_init(),
activation_fn=None,
scope='input_projection',
reuse=reuse)
for i in range(num_layers):
T = fc(x,
size,
weights_initializer=xa_init(),
biases_initializer=bias_init(),
activation_fn=tf.sigmoid,
scope='gate_%d' % i,
reuse=reuse)
H = fc(x,
size,
weights_initializer=he_init(),
biases_initializer=bias_init(),
activation_fn=activation,
scope='activation_%d' % i,
reuse=reuse)
H = tf.nn.dropout(H, 1.0 - dropout)
x = H * T + x * (1.0 - T)
return x
def fc(x,
channels,
normal_weight_init=False,
activation_fn='leaky',
scope='fc_0'):
with tf.variable_scope(scope):
x = tf.layers.flatten(x)
if normal_weight_init:
x = tf.layers.dense(
inputs=x,
units=channels,
kernel_initializer=tf.truncated_normal_initializer(
stddev=0.02),
kernel_regularizer=tf_contrib.layers.l2_regularizer(
scale=0.0001))
else:
if activation_fn == 'relu':
x = tf.layers.dense(
inputs=x,
units=channels,
kernel_initializer=he_init(),
kernel_regularizer=tf_contrib.layers.l2_regularizer(
scale=0.0001))
else:
x = tf.layers.dense(
inputs=x,
units=channels,
kernel_size=kernel,
kernel_regularizer=tf_contrib.layers.l2_regularizer(
scale=0.0001))
x = activation(x, activation_fn)
return x
def encoder_block(inputs,
num_conv_layers,
kernel_size,
hidden_size,
num_heads,
num_blocks=1,
mask=None,
dropout=0.0,
use_relative_pos=False,
scope='encoder_block',
reuse=None):
with tf.variable_scope(scope, reuse=reuse):
outputs = inputs
for block_repeat_idx in range(num_blocks):
with tf.variable_scope('block_%d' % block_repeat_idx, reuse=reuse):
total_sublayers = num_conv_layers + 2
sublayer = 0
# position encoding
if not use_relative_pos:
outputs += get_timing_signal_1d(
tf.shape(outputs)[1],
tf.shape(outputs)[2])
# convolutions
outputs = tf.expand_dims(outputs, 2)
for i in range(num_conv_layers):
residual = outputs
outputs = layer_norm(
outputs,
begin_norm_axis=-1,
begin_params_axis=-1,
scope='conv_layer_norm_%d' % i,
reuse=reuse) # TODO: change layer norm
if i % 2 == 0:
outputs = tf.nn.dropout(outputs, 1.0 - dropout)
if isinstance(kernel_size, list):
kernel_num = len(kernel_size)
kernel_outputs = [
depthwise_conv(
outputs,
hidden_size,
bias=True,
activation=tf.nn.relu,
kernel_size=k,
scope='depthwise_conv_%d_kernel_%d' % (i, k),
reuse=reuse) for k in kernel_size
]
kernel_weights = tf.nn.softmax(tf.get_variable(
'kernel_weights_conv_%d' % i, [kernel_num],
dtype=tf.float32,
trainable=True,
initializer=tf.constant_initializer(1.0 /
kernel_num)),
axis=0)
outputs = 0
for j in range(kernel_num):
outputs += kernel_outputs[j] * kernel_weights[j]
else:
outputs = depthwise_conv(
outputs,
hidden_size,
bias=True,
activation=tf.nn.relu,
# activation=tf.nn.relu if i < num_conv_layers - 1 else None,
kernel_size=kernel_size,
scope='depthwise_conv_%d' % i,
reuse=reuse)
sublayer += 1
outputs = layer_dropout(
residual,
residual + tf.nn.dropout(outputs, 1 - dropout),
dropout * float(sublayer) / total_sublayers)
outputs = tf.squeeze(outputs, 2)
# self attention
residual = outputs
outputs = layer_norm(outputs,
begin_norm_axis=-1,
begin_params_axis=-1,
scope='self_attention_layer_norm',
reuse=reuse)
outputs = tf.nn.dropout(outputs, 1.0 - dropout)
outputs = self_attention(outputs,
hidden_size,
num_heads,
use_relative_pos=use_relative_pos,
mask=mask,
scope='self_attention_layer',
reuse=reuse)
sublayer += 1
outputs = layer_dropout(
residual, residual + tf.nn.dropout(outputs, 1 - dropout),
dropout * float(sublayer) / total_sublayers)
# feed forward
residual = outputs
outputs = layer_norm(outputs,
begin_norm_axis=-1,
begin_params_axis=-1,
scope='fc_layer_norm',
reuse=reuse)
outputs = tf.nn.dropout(outputs, 1.0 - dropout)
outputs = fc(outputs,
hidden_size,
tf.nn.relu,
weights_initializer=he_init(),
biases_initializer=bias_init(),
scope='fc_layer_1',
reuse=reuse)
# outputs = tf.nn.dropout(outputs, 1 - dropout)
outputs = fc(outputs,
hidden_size,
None,
weights_initializer=xa_init(),
biases_initializer=bias_init(),
scope='fc_layer_2',
reuse=reuse)
sublayer += 1
outputs = layer_dropout(
residual, residual + tf.nn.dropout(outputs, 1 - dropout),
dropout * float(sublayer) / total_sublayers)
return outputs
def conv_layer(x, channels, kernel=3, stride=1, activation='relu', padding='VALID', layer_name='_conv3d') :
with tf.name_scope(layer_name) :
if activation == None :
return tf.layers.conv3d(inputs=x, filters=channels, kernel_size=kernel, kernel_initializer=he_init(),
strides=stride, padding=padding)
else :
return tf.layers.conv3d(inputs=x, filters=channels, kernel_size=kernel, kernel_initializer=he_init(),
strides=stride, padding=padding, activation=relu)
def deconv_layer(x, channels, kernel=4, stride=2, padding='SAME', layer_name='_deconv3d') :
with tf.name_scope(layer_name) :
x = tf.layers.conv3d_transpose(inputs=x, filters=channels, kernel_size=kernel, kernel_initializer=he_init(),
strides=stride, padding=padding, use_bias=False)
return x
def conv_layer(x, channels, kernel=3, stride=1, activation='relu', padding='VALID', layer_name='_conv3d') :
with tf.name_scope(layer_name) :
if activation == None :
return tf.layers.conv3d(inputs=x, filters=channels, kernel_size=kernel, kernel_initializer=he_init(),
strides=stride, padding=padding)
else :
return tf.layers.conv3d(inputs=x, filters=channels, kernel_size=kernel, kernel_initializer=he_init(),
strides=stride, padding=padding, activation=relu)
def deconv_layer(x, channels, kernel=4, stride=2, padding='VALID', layer_name='_deconv3d') :
with tf.name_scope(layer_name) :
crop = 1
x = tf.layers.conv3d_transpose(inputs=x, filters=channels, kernel_size=kernel, kernel_initializer=he_init(),
strides=stride, padding=padding, use_bias=False)
x = x[:, crop:-crop, crop:-crop, crop:-crop, :]
return x
