def concat(num_ins):
def compile_fn(di, dh):
def forward_fn(di, is_training=True):
return {
'out':
tf.concat(values=[di[input_name] for input_name in di], axis=3)
}
return forward_fn
return htfe.TensorflowEagerModule('Concat', compile_fn, {},
['in' + str(i) for i in range(num_ins)],
['out']).get_io()
def drop_path(cell_ratio):
def compile_fn(di, dh):
drop_path_layer = DropPath(dh['keep_prob'], cell_ratio, di['in1'])
def forward_fn(di, is_training=True):
return {'out': drop_path_layer(di['in0'])}
return forward_fn
return htfe.TensorflowEagerModule(
'DropPath', compile_fn, {
'keep_prob': hp_sharer.get('drop_path_keep_prob')
}, ['in0', 'in1'], ['out']).get_io()
def add(num_inputs):
def compile_fn(di, dh):
in_channels = [tf.shape(di[inp])[-1] for inp in di]
final_channels = tf.reduce_min(in_channels)
def forward_fn(di, is_training=True):
out = [di[inp][:, :, :, :final_channels] for inp in di]
return {'out': tf.add_n(out)}
return forward_fn
return htfe.TensorflowEagerModule(
'Add', compile_fn, {}, ['in' + str(i) for i in range(num_inputs)],
['out']).get_io()
def concat(num_ins):
def compile_fn(di, dh):
concat_op = tf.keras.layers.Concatenate(
axis=3) if num_ins > 1 else None
def forward_fn(di, is_training=True):
return {
'out':
concat_op([di[input_name]
for input_name in di]) if concat_op else di['in0']
}
return forward_fn
return htfe.TensorflowEagerModule('Concat', compile_fn, {},
['in' + str(i) for i in range(num_ins)],
['out']).get_io()
def concat(num_inputs):
input_names = ["in%d" % i for i in range(num_inputs)]
def compile_fn(di, dh):
if num_inputs > 1:
concat = Concatenate()
def forward_fn(di, is_training=False):
return {
"out":
concat([v for name, v in di.items()])
if num_inputs > 1 else di['in0']
}
return forward_fn
return htfe.TensorflowEagerModule("ConcatCombiner", compile_fn, {},
input_names, ['out']).get_io()
def maybe_factorized_reduction(add_relu=False):
def compile_fn(di, dh):
_, _, height, channels = di['in0'].shape.as_list()
_, _, final_height, final_channels = di['in1'].shape.as_list()
if add_relu:
relu = tf.keras.layers.ReLU()
if height == final_height and channels == final_channels:
pass
elif height == final_height:
conv = tf.keras.layers.Conv2D(final_channels, 1)
bn = tf.keras.layers.BatchNormalization()
else:
avg_pool = tf.keras.layers.AveragePooling2D(1, 2, padding='VALID')
conv1 = tf.keras.layers.Conv2D(int(final_channels / 2), 1)
conv2 = tf.keras.layers.Conv2D(int(final_channels / 2), 1)
bn = tf.keras.layers.BatchNormalization()
pad = tf.keras.layers.ZeroPadding2D(((0, 1), (0, 1)))
slice_layer = tf.keras.layers.Lambda(lambda x: x[:, 1:, 1:, :])
concat = tf.keras.layers.Concatenate(axis=3)
def forward_fn(di, is_training=True):
inp = relu(di['in0']) if add_relu else di['in0']
if height == final_height and channels == final_channels:
out = inp
elif height == final_height:
out = conv(inp)
out = bn(out)
else:
path1 = avg_pool(inp)
path1 = conv1(path1)
path2 = slice_layer(pad(inp))
path2 = avg_pool(path2)
path2 = conv2(path2)
out = concat([path1, path2])
out = bn(out)
return {'out': out}
return forward_fn
return htfe.TensorflowEagerModule('MaybeFactorizedReduction', compile_fn,
{}, ['in0', 'in1'], ['out']).get_io()