def conv2d(h_num_filters,
h_filter_width,
h_stride=1,
h_dilation_rate=1,
h_use_bias=True,
h_padding='SAME'):
def compile_fn(di, dh):
conv = tf.keras.layers.Conv2D(dh['num_filters'],
dh['filter_width'],
dh['stride'],
use_bias=dh['use_bias'],
dilation_rate=dh['dilation_rate'],
padding=dh['padding'])
def forward_fn(di, is_training=True):
return {'out': conv(di['in'])}
return forward_fn
return siso_tensorflow_eager_module(
'Conv2D', compile_fn, {
'num_filters': h_num_filters,
'filter_width': h_filter_width,
'stride': h_stride,
'dilation_rate': h_dilation_rate,
'use_bias': h_use_bias,
'padding': h_padding
})
def separable_conv2d(h_num_filters,
h_filter_width,
h_stride=1,
h_dilation_rate=1,
h_depth_multiplier=1,
h_use_bias=True,
h_padding='SAME'):
def compile_fn(di, dh):
conv_op = tf.keras.layers.SeparableConv2D(
dh['num_filters'],
dh['filter_width'],
strides=dh['stride'],
dilation_rate=dh['dilation_rate'],
depth_multiplier=dh['depth_multiplier'],
use_bias=dh['use_bias'],
padding=dh['padding'])
def fn(di, is_training=True):
return {'out': conv_op(di['in'])}
return fn
return siso_tensorflow_eager_module(
'SeparableConv2D', compile_fn, {
'num_filters': h_num_filters,
'filter_width': h_filter_width,
'stride': h_stride,
'use_bias': h_use_bias,
'dilation_rate': h_dilation_rate,
'depth_multiplier': h_depth_multiplier,
'padding': h_padding
})
def relu():
def compile_fn(di, dh):
def forward_fn(di, is_training=True):
return {'out': tf.nn.relu(di['in'])}
return forward_fn
return siso_tensorflow_eager_module('ReLU', compile_fn, {})
def flatten():
def compile_fn(di, dh):
def forward_fn(di, is_training=True):
return {'out': tf.layers.flatten(di['in'])}
return forward_fn
return siso_tensorflow_eager_module('Flatten', compile_fn, {})
def global_pool2d():
def compile_fn(di, dh):
def forward_fn(di, is_training=True):
return {'out': tf.reduce_mean(di['in'], [1, 2])}
return forward_fn
return siso_tensorflow_eager_module('GlobalAveragePool', compile_fn, {})
def global_pool():
def compile_fn(di, dh):
def fn(di, is_training):
with tf.device('/gpu:0'):
return {'out': tf.reduce_mean(di['in'], [1, 2])}
return fn
return htfe.siso_tensorflow_eager_module('GlobalPool', compile_fn, {})
def relu():
def compile_fn(di, dh):
def fn(di, is_training=True):
with tf.device('/gpu:0'):
return {'out': tf.nn.relu(di['in'])}
return fn
return htfe.siso_tensorflow_eager_module('ReLU', compile_fn, {})
def batch_normalization():
def compile_fn(di, dh):
bn = tf.layers.BatchNormalization(momentum=.9, epsilon=1e-5)
def forward_fn(di, is_training):
return {'out': bn(di['in'], training=is_training)}
return forward_fn
return siso_tensorflow_eager_module('BatchNormalization', compile_fn, {})
def global_pool2d():
def compile_fn(di, dh):
pool = tf.keras.layers.GlobalAveragePooling2D()
def forward_fn(di, is_training=True):
return {'out': pool(di['in'])}
return forward_fn
return siso_tensorflow_eager_module('GlobalAveragePool', compile_fn, {})
def fc_layer(h_num_units):
def compile_fn(di, dh):
fc = tf.layers.Dense(dh['num_units'])
def forward_fn(di, is_training=True):
return {'out': fc(di['in'])}
return forward_fn
return siso_tensorflow_eager_module('FCLayer', compile_fn,
{'num_units': h_num_units})
def conv2D(filter_size, channels, name):
def compile_fn(di, dh):
conv_fn = lambda: tf.keras.layers.Conv2D(channels, filter_size)
conv = weight_sharer.get(name, conv_fn)
def fn(di, is_training=True):
return {'out': conv(di['in'])}
return fn
return siso_tensorflow_eager_module('Conv2D', compile_fn, {})
def dropout(keep_prob):
def compile_fn(di, dh):
def fn(di, is_training=True):
if is_training:
with tf.device('/gpu:0'):
out = tf.nn.dropout(di['in'], keep_prob)
else:
out = di['in']
return {'out': out}
return fn
return htfe.siso_tensorflow_eager_module('Dropout', compile_fn, {})
def dropout(h_keep_prob):
def compile_fn(di, dh):
def forward_fn(di, is_training=True):
if is_training:
out = tf.nn.dropout(di['in'], dh['keep_prob'])
else:
out = di['in']
return {'out': out}
return forward_fn
return siso_tensorflow_eager_module('Dropout', compile_fn,
{'keep_prob': h_keep_prob})
def global_convolution(h_num_filters):
def compile_fn(di, dh):
_, h, w, _ = di['in'].shape.as_list()
conv = tf.keras.layers.Conv2D(dh['num_filters'], [h, w],
use_bias=False,
padding='VALID')
def forward_fn(di, is_training=True):
return {'out': conv(di['in'])}
return forward_fn
return htfe.siso_tensorflow_eager_module('GlobalConv2D', compile_fn, {
'num_filters': h_num_filters,
})
def max_pool2d(h_kernel_size, h_stride=1, h_padding='SAME'):
def compile_fn(di, dh):
pool = tf.layers.MaxPooling2D(dh['kernel_size'],
dh['stride'],
padding=dh['padding'])
def forward_fn(di, is_training=True):
return {'out': pool(di['in'])}
return forward_fn
return siso_tensorflow_eager_module('MaxPool2D', compile_fn, {
'kernel_size': h_kernel_size,
'stride': h_stride,
'padding': h_padding
})
def check_filters(filters, stride=1):
def compile_fn(di, dh):
num_filters = di['in'].shape[-1].value
if num_filters != filters or stride > 1:
conv = tf.keras.layers.Conv2D(filters,
1,
strides=stride,
padding='SAME')
else:
conv = None
def forward_fn(di, is_training=True):
return {'out': di['in'] if conv is None else conv(di['in'])}
return forward_fn
return htfe.siso_tensorflow_eager_module('CheckFilters', compile_fn, {})
def avg_pool(h_kernel_size, h_stride):
def compile_fn(di, dh):
def fn(di, is_training=True):
with tf.device('/gpu:0'):
return {
'out':
tf.nn.avg_pool(
di['in'], [1, dh['kernel_size'], dh['kernel_size'], 1],
[1, dh['stride'], dh['stride'], 1], 'SAME')
}
return fn
return htfe.siso_tensorflow_eager_module('AvgPool', compile_fn, {
'kernel_size': h_kernel_size,
'stride': h_stride,
})
def min_pool2d(h_kernel_size, h_stride=1, h_padding='SAME'):
def compile_fn(di, dh):
pool = tf.keras.layers.MaxPooling2D(dh['kernel_size'],
dh['stride'],
padding=dh['padding'])
negate = tf.keras.layers.Lambda(lambda x: -1 * x, name='negate')
def forward_fn(di, is_training=True):
return {'out': negate(pool(negate(di['in'])))}
return forward_fn
return siso_tensorflow_eager_module('MinPool2D', compile_fn, {
'kernel_size': h_kernel_size,
'stride': h_stride,
'padding': h_padding
})
def dropout(h_keep_prob):
class Dropout(tf.keras.layers.Layer):
def __init__(self, rate, seed=None, **kwargs):
super(Dropout, self).__init__(**kwargs)
self.rate = rate
self.seed = seed
self.supports_masking = True
def call(self, inputs, training=None):
if training is None:
training = True
def dropped_inputs():
return tf.nn.dropout(inputs, 1 - self.rate, seed=self.seed)
output = tf_utils.smart_cond(
training, dropped_inputs, lambda: array_ops.identity(inputs))
return output
def compute_output_shape(self, input_shape):
return input_shape
def get_config(self):
config = {'rate': self.rate, 'seed': self.seed}
base_config = super(Dropout, self).get_config()
return dict(list(base_config.items()) + list(config.items()))
def compile_fn(di, dh):
dropout_op = Dropout(1 - dh['keep_prob'])
def forward_fn(di, is_training=True):
# out = tf.nn.dropout(di['in'], dh['keep_prob'])
# else:
# out = di['in']
out = dropout_op(di['in'])
return {'out': out}
return forward_fn
return siso_tensorflow_eager_module('Dropout', compile_fn,
{'keep_prob': h_keep_prob})
def keras_batch_normalization(name='default', weight_sharer=None):
name = name + '_bn'
def compile_fn(di, dh):
bn = weight_sharer.get(name, tf.keras.layers.BatchNormalization,
lambda layer: layer.get_weights())
if not bn.built:
with tf.device('/gpu:0'):
bn.build(di['in'].get_shape())
weights = weight_sharer.load_weights(name)
if weights is not None:
bn.set_weights(weights)
def fn(di, is_training):
with tf.device('/gpu:0'):
return {'out': bn(di['in'], training=is_training)}
return fn
return htfe.siso_tensorflow_eager_module('BatchNormalization', compile_fn,
{})
def fc_layer(num_classes, name, weight_sharer):
name = name + '_fc_layer_' + str(num_classes)
def compile_fn(di, dh):
fc = weight_sharer.get(name,
lambda: tf.keras.layers.Dense(num_classes),
lambda layer: layer.get_weights())
if not fc.built:
with tf.device('/gpu:0'):
fc.build(di['in'].get_shape())
weights = weight_sharer.load_weights(name)
if weights is not None:
fc.set_weights(weights)
def fn(di, is_training=True):
with tf.device('/gpu:0'):
return {'out': fc(di['in'])}
return fn
return htfe.siso_tensorflow_eager_module('FC_Layer', compile_fn, {})
def conv2D_depth_separable(filter_size, name, weight_sharer, out_filters=None):
def compile_fn(di, dh):
(_, _, _, channels) = di['in'].get_shape().as_list()
channels = channels if out_filters is None else out_filters
conv_fn = lambda: tf.keras.layers.SeparableConv2D(
channels, filter_size, padding='same')
conv = weight_sharer.get(name + '_dsep_' + str(filter_size), conv_fn,
lambda layer: layer.get_weights())
if not conv.built:
with tf.device('/gpu:0'):
conv.build(di['in'].get_shape())
weights = weight_sharer.load_weights(name)
if weights is not None:
conv.set_weights(weights)
def fn(di, is_training=True):
with tf.device('/gpu:0'):
return {'out': conv(di['in'])}
return fn
return htfe.siso_tensorflow_eager_module('Conv2DSeparable', compile_fn, {})
