# -*- encoding:utf-8 -*- # Copyright (c) Alibaba, Inc. and its affiliates. import logging import six import tensorflow as tf from google.protobuf import struct_pb2 from easy_rec.python.layers.common_layers import EnhancedInputLayer from easy_rec.python.layers.keras import MLP from easy_rec.python.layers.keras import EmbeddingLayer from easy_rec.python.layers.utils import Parameter from easy_rec.python.protos import backbone_pb2 from easy_rec.python.utils.dag import DAG from easy_rec.python.utils.load_class import load_keras_layer from easy_rec.python.utils.tf_utils import add_elements_to_collection if tf.__version__ >= '2.0': tf = tf.compat.v1 class Package(object): """A sub DAG of tf ops for reuse.""" __packages = {} @staticmethod def has_backbone_block(name): if 'backbone' not in Package.__packages: return False backbone = Package.__packages['backbone'] return backbone.has_block(name) @staticmethod def backbone_block_outputs(name): if 'backbone' not in Package.__packages: return None backbone = Package.__packages['backbone'] return backbone.block_outputs(name) def __init__(self, config, features, input_layer, l2_reg=None): self._config = config self._features = features self._input_layer = input_layer self._l2_reg = l2_reg self._dag = DAG() self._name_to_blocks = {} self._name_to_layer = {} self.reset_input_config(None) self._block_outputs = {} self._package_input = None self._feature_group_inputs = {} reuse = None if config.name == 'backbone' else tf.AUTO_REUSE input_feature_groups = self._feature_group_inputs for block in config.blocks: if len(block.inputs) == 0: raise ValueError('block takes at least one input: %s' % block.name) self._dag.add_node(block.name) self._name_to_blocks[block.name] = block layer = block.WhichOneof('layer') if layer in {'input_layer', 'raw_input', 'embedding_layer'}: if len(block.inputs) != 1: raise ValueError('input layer `%s` takes only one input' % block.name) one_input = block.inputs[0] name = one_input.WhichOneof('name') if name != 'feature_group_name': raise KeyError( '`feature_group_name` should be set for input layer: ' + block.name) group = one_input.feature_group_name if not input_layer.has_group(group): raise KeyError('invalid feature group name: ' + group) if group in input_feature_groups: if layer == input_layer: logging.warning('input `%s` already exists in other block' % group) elif layer == 'raw_input': input_fn = input_feature_groups[group] self._name_to_layer[block.name] = input_fn elif layer == 'embedding_layer': inputs, vocab, weights = input_feature_groups[group] block.embedding_layer.vocab_size = vocab params = Parameter.make_from_pb(block.embedding_layer) input_fn = EmbeddingLayer(params, block.name) self._name_to_layer[block.name] = input_fn else: if layer == 'input_layer': input_fn = EnhancedInputLayer(self._input_layer, self._features, group, reuse) input_feature_groups[group] = input_fn elif layer == 'raw_input': input_fn = self._input_layer.get_raw_features(self._features, group) input_feature_groups[group] = input_fn else: # embedding_layer inputs, vocab, weights = self._input_layer.get_bucketized_features( self._features, group) block.embedding_layer.vocab_size = vocab params = Parameter.make_from_pb(block.embedding_layer) input_fn = EmbeddingLayer(params, block.name) input_feature_groups[group] = (inputs, vocab, weights) logging.info('add an embedding layer %s with vocab size %d', block.name, vocab) self._name_to_layer[block.name] = input_fn else: self.define_layers(layer, block, block.name, reuse) # sequential layers for i, layer_cnf in enumerate(block.layers): layer = layer_cnf.WhichOneof('layer') name_i = '%s_l%d' % (block.name, i) self.define_layers(layer, layer_cnf, name_i, reuse) num_groups = len(input_feature_groups) num_blocks = len(self._name_to_blocks) - num_groups assert num_blocks > 0, 'there must be at least one block in backbone' num_pkg_input = 0 for block in config.blocks: layer = block.WhichOneof('layer') if layer in {'input_layer', 'raw_input', 'embedding_layer'}: continue name = block.name if name in input_feature_groups: raise KeyError('block name can not be one of feature groups:' + name) for input_node in block.inputs: input_type = input_node.WhichOneof('name') input_name = getattr(input_node, input_type) if input_type == 'use_package_input': assert input_name, 'use_package_input can not set false' num_pkg_input += 1 continue if input_type == 'package_name': num_pkg_input += 1 self._dag.add_node_if_not_exists(input_name) self._dag.add_edge(input_name, name) if input_node.HasField('package_input'): pkg_input_name = input_node.package_input self._dag.add_node_if_not_exists(pkg_input_name) self._dag.add_edge(pkg_input_name, input_name) continue iname = input_name if iname in self._name_to_blocks: assert iname != name, 'input name can not equal to block name:' + iname self._dag.add_edge(iname, name) else: is_fea_group = input_type == 'feature_group_name' if is_fea_group and input_layer.has_group(iname): logging.info('adding an input_layer block: ' + iname) new_block = backbone_pb2.Block() new_block.name = iname input_cfg = backbone_pb2.Input() input_cfg.feature_group_name = iname new_block.inputs.append(input_cfg) new_block.input_layer.CopyFrom(backbone_pb2.InputLayer()) self._name_to_blocks[iname] = new_block self._dag.add_node(iname) self._dag.add_edge(iname, name) if iname in input_feature_groups: fn = input_feature_groups[iname] else: fn = EnhancedInputLayer(self._input_layer, self._features, iname) input_feature_groups[iname] = fn self._name_to_layer[iname] = fn elif Package.has_backbone_block(iname): backbone = Package.__packages['backbone'] backbone._dag.add_node_if_not_exists(self._config.name) backbone._dag.add_edge(iname, self._config.name) num_pkg_input += 1 else: raise KeyError( 'invalid input name `%s`, must be the name of either a feature group or an another block' % iname) num_groups = len(input_feature_groups) assert num_pkg_input > 0 or num_groups > 0, 'there must be at least one input layer/feature group' if len(config.concat_blocks) == 0 and len(config.output_blocks) == 0: leaf = self._dag.all_leaves() logging.warning( '%s has no `concat_blocks` or `output_blocks`, try to concat all leaf blocks: %s' % (config.name, ','.join(leaf))) self._config.concat_blocks.extend(leaf) Package.__packages[self._config.name] = self logging.info('%s layers: %s' % (config.name, ','.join(self._name_to_layer.keys()))) def define_layers(self, layer, layer_cnf, name, reuse): if layer == 'keras_layer': layer_obj = self.load_keras_layer(layer_cnf.keras_layer, name, reuse) self._name_to_layer[name] = layer_obj elif layer == 'recurrent': keras_layer = layer_cnf.recurrent.keras_layer for i in range(layer_cnf.recurrent.num_steps): name_i = '%s_%d' % (name, i) layer_obj = self.load_keras_layer(keras_layer, name_i, reuse) self._name_to_layer[name_i] = layer_obj elif layer == 'repeat': keras_layer = layer_cnf.repeat.keras_layer for i in range(layer_cnf.repeat.num_repeat): name_i = '%s_%d' % (name, i) layer_obj = self.load_keras_layer(keras_layer, name_i, reuse) self._name_to_layer[name_i] = layer_obj def reset_input_config(self, config): self.input_config = config def set_package_input(self, pkg_input): self._package_input = pkg_input def has_block(self, name): return name in self._name_to_blocks def block_outputs(self, name): return self._block_outputs.get(name, None) def block_input(self, config, block_outputs, training=None, **kwargs): inputs = [] for input_node in config.inputs: input_type = input_node.WhichOneof('name') input_name = getattr(input_node, input_type) if input_type == 'use_package_input': input_feature = self._package_input input_name = 'package_input' elif input_type == 'package_name': if input_name not in Package.__packages: raise KeyError('package name `%s` does not exists' % input_name) package = Package.__packages[input_name] if input_node.HasField('reset_input'): package.reset_input_config(input_node.reset_input) if input_node.HasField('package_input'): pkg_input_name = input_node.package_input if pkg_input_name in block_outputs: pkg_input = block_outputs[pkg_input_name] else: if pkg_input_name not in Package.__packages: raise KeyError('package name `%s` does not exists' % pkg_input_name) inner_package = Package.__packages[pkg_input_name] pkg_input = inner_package(training) if input_node.HasField('package_input_fn'): fn = eval(input_node.package_input_fn) pkg_input = fn(pkg_input) package.set_package_input(pkg_input) input_feature = package(training, **kwargs) elif input_name in block_outputs: input_feature = block_outputs[input_name] else: input_feature = Package.backbone_block_outputs(input_name) if input_feature is None: raise KeyError('input name `%s` does not exists' % input_name) if input_node.ignore_input: continue if input_node.HasField('input_slice'): fn = eval('lambda x: x' + input_node.input_slice.strip()) input_feature = fn(input_feature) if input_node.HasField('input_fn'): with tf.name_scope(config.name): fn = eval(input_node.input_fn) input_feature = fn(input_feature) inputs.append(input_feature) if config.merge_inputs_into_list: output = inputs else: try: output = merge_inputs(inputs, config.input_concat_axis, config.name) except ValueError as e: msg = getattr(e, 'message', str(e)) logging.error('merge inputs of block %s failed: %s', config.name, msg) raise e if config.HasField('extra_input_fn'): fn = eval(config.extra_input_fn) output = fn(output) return output def __call__(self, is_training, **kwargs): with tf.name_scope(self._config.name): return self.call(is_training, **kwargs) def call(self, is_training, **kwargs): block_outputs = {} self._block_outputs = block_outputs # reset blocks = self._dag.topological_sort() logging.info(self._config.name + ' topological order: ' + ','.join(blocks)) for block in blocks: if block not in self._name_to_blocks: assert block in Package.__packages, 'invalid block: ' + block continue config = self._name_to_blocks[block] if config.layers: # sequential layers logging.info('call sequential %d layers' % len(config.layers)) output = self.block_input(config, block_outputs, is_training, **kwargs) for i, layer in enumerate(config.layers): name_i = '%s_l%d' % (block, i) output = self.call_layer(output, layer, name_i, is_training, **kwargs) block_outputs[block] = output continue # just one of layer layer = config.WhichOneof('layer') if layer is None: # identity layer output = self.block_input(config, block_outputs, is_training, **kwargs) block_outputs[block] = output elif layer == 'raw_input': block_outputs[block] = self._name_to_layer[block] elif layer == 'input_layer': input_fn = self._name_to_layer[block] input_config = config.input_layer if self.input_config is not None: input_config = self.input_config input_fn.reset(input_config, is_training) block_outputs[block] = input_fn(input_config, is_training) elif layer == 'embedding_layer': input_fn = self._name_to_layer[block] feature_group = config.inputs[0].feature_group_name inputs, _, weights = self._feature_group_inputs[feature_group] block_outputs[block] = input_fn([inputs, weights], is_training) else: with tf.name_scope(block + '_input'): inputs = self.block_input(config, block_outputs, is_training, **kwargs) output = self.call_layer(inputs, config, block, is_training, **kwargs) block_outputs[block] = output outputs = [] for output in self._config.output_blocks: if output in block_outputs: temp = block_outputs[output] outputs.append(temp) else: raise ValueError('No output `%s` of backbone to be concat' % output) if outputs: return outputs for output in self._config.concat_blocks: if output in block_outputs: temp = block_outputs[output] outputs.append(temp) else: raise ValueError('No output `%s` of backbone to be concat' % output) try: output = merge_inputs(outputs, msg='backbone') except ValueError as e: msg = getattr(e, 'message', str(e)) logging.error("merge backbone's output failed: %s", msg) raise e return output def load_keras_layer(self, layer_conf, name, reuse=None): layer_cls, customize = load_keras_layer(layer_conf.class_name) if layer_cls is None: raise ValueError('Invalid keras layer class name: ' + layer_conf.class_name) param_type = layer_conf.WhichOneof('params') if customize: if param_type is None or param_type == 'st_params': params = Parameter(layer_conf.st_params, True, l2_reg=self._l2_reg) else: pb_params = getattr(layer_conf, param_type) params = Parameter(pb_params, False, l2_reg=self._l2_reg) has_reuse = True try: from funcsigs import signature sig = signature(layer_cls.__init__) has_reuse = 'reuse' in sig.parameters.keys() except ImportError: try: from sklearn.externals.funcsigs import signature sig = signature(layer_cls.__init__) has_reuse = 'reuse' in sig.parameters.keys() except ImportError: logging.warning('import funcsigs failed') if has_reuse: layer = layer_cls(params, name=name, reuse=reuse) else: layer = layer_cls(params, name=name) return layer, customize elif param_type is None: # internal keras layer layer = layer_cls(name=name) return layer, customize else: assert param_type == 'st_params', 'internal keras layer only support st_params' try: kwargs = convert_to_dict(layer_conf.st_params) logging.info('call %s layer with params %r' % (layer_conf.class_name, kwargs)) layer = layer_cls(name=name, **kwargs) except TypeError as e: logging.warning(e) args = map(format_value, layer_conf.st_params.values()) logging.info('try to call %s layer with params %r' % (layer_conf.class_name, args)) layer = layer_cls(*args, name=name) return layer, customize def call_keras_layer(self, inputs, name, training, **kwargs): """Call predefined Keras Layer, which can be reused.""" layer, customize = self._name_to_layer[name] cls = layer.__class__.__name__ if customize: try: output = layer(inputs, training=training, **kwargs) except Exception as e: msg = getattr(e, 'message', str(e)) logging.error('call keras layer %s (%s) failed: %s' % (name, cls, msg)) raise e else: try: output = layer(inputs, training=training) if cls == 'BatchNormalization': add_elements_to_collection(layer.updates, tf.GraphKeys.UPDATE_OPS) except TypeError: output = layer(inputs) return output def call_layer(self, inputs, config, name, training, **kwargs): layer_name = config.WhichOneof('layer') if layer_name == 'keras_layer': return self.call_keras_layer(inputs, name, training, **kwargs) if layer_name == 'lambda': conf = getattr(config, 'lambda') fn = eval(conf.expression) return fn(inputs) if layer_name == 'repeat': conf = config.repeat n_loop = conf.num_repeat outputs = [] for i in range(n_loop): name_i = '%s_%d' % (name, i) ly_inputs = inputs if conf.HasField('input_slice'): fn = eval('lambda x, i: x' + conf.input_slice.strip()) ly_inputs = fn(ly_inputs, i) if conf.HasField('input_fn'): with tf.name_scope(config.name): fn = eval(conf.input_fn) ly_inputs = fn(ly_inputs, i) output = self.call_keras_layer(ly_inputs, name_i, training, **kwargs) outputs.append(output) if len(outputs) == 1: return outputs[0] if conf.HasField('output_concat_axis'): return tf.concat(outputs, conf.output_concat_axis) return outputs if layer_name == 'recurrent': conf = config.recurrent fixed_input_index = -1 if conf.HasField('fixed_input_index'): fixed_input_index = conf.fixed_input_index if fixed_input_index >= 0: assert type(inputs) in (tuple, list), '%s inputs must be a list' output = inputs for i in range(conf.num_steps): name_i = '%s_%d' % (name, i) output_i = self.call_keras_layer(output, name_i, training, **kwargs) if fixed_input_index >= 0: j = 0 for idx in range(len(output)): if idx == fixed_input_index: continue if type(output_i) in (tuple, list): output[idx] = output_i[j] else: output[idx] = output_i j += 1 else: output = output_i if fixed_input_index >= 0: del output[fixed_input_index] if len(output) == 1: return output[0] return output return output raise NotImplementedError('Unsupported backbone layer:' + layer_name) class Backbone(object): """Configurable Backbone Network.""" def __init__(self, config, features, input_layer, l2_reg=None): self._config = config self._l2_reg = l2_reg main_pkg = backbone_pb2.BlockPackage() main_pkg.name = 'backbone' main_pkg.blocks.MergeFrom(config.blocks) if config.concat_blocks: main_pkg.concat_blocks.extend(config.concat_blocks) if config.output_blocks: main_pkg.output_blocks.extend(config.output_blocks) self._main_pkg = Package(main_pkg, features, input_layer, l2_reg) for pkg in config.packages: Package(pkg, features, input_layer, l2_reg) def __call__(self, is_training, **kwargs): output = self._main_pkg(is_training, **kwargs) if self._config.HasField('top_mlp'): params = Parameter.make_from_pb(self._config.top_mlp) params.l2_regularizer = self._l2_reg final_mlp = MLP(params, name='backbone_top_mlp') if type(output) in (list, tuple): output = tf.concat(output, axis=-1) output = final_mlp(output, training=is_training, **kwargs) return output @classmethod def wide_embed_dim(cls, config): wide_embed_dim = None for pkg in config.packages: wide_embed_dim = get_wide_embed_dim(pkg.blocks, wide_embed_dim) return get_wide_embed_dim(config.blocks, wide_embed_dim) def get_wide_embed_dim(blocks, wide_embed_dim=None): for block in blocks: layer = block.WhichOneof('layer') if layer == 'input_layer': if block.input_layer.HasField('wide_output_dim'): wide_dim = block.input_layer.wide_output_dim if wide_embed_dim: assert wide_embed_dim == wide_dim, 'wide_output_dim must be consistent' else: wide_embed_dim = wide_dim return wide_embed_dim def merge_inputs(inputs, axis=-1, msg=''): if len(inputs) == 0: raise ValueError('no inputs to be concat:' + msg) if len(inputs) == 1: return inputs[0] from functools import reduce if all(map(lambda x: type(x) == list, inputs)): # merge multiple lists into a list return reduce(lambda x, y: x + y, inputs) if any(map(lambda x: type(x) == list, inputs)): logging.warning('%s: try to merge inputs into list' % msg) return reduce(lambda x, y: x + y, [e if type(e) == list else [e] for e in inputs]) if axis != -1: logging.info('concat inputs %s axis=%d' % (msg, axis)) return tf.concat(inputs, axis=axis) def format_value(value): value_type = type(value) if value_type == six.text_type: return str(value) if value_type == float: int_v = int(value) return int_v if int_v == value else value if value_type == struct_pb2.ListValue: return map(format_value, value) if value_type == struct_pb2.Struct: return convert_to_dict(value) return value def convert_to_dict(struct): kwargs = {} for key, value in struct.items(): kwargs[str(key)] = format_value(value) return kwargs