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()