# -*- encoding:utf-8 -*- # Copyright (c) Alibaba, Inc. and its affiliates. from __future__ import absolute_import from __future__ import division from __future__ import print_function import json import logging import math import os import time import six import tensorflow as tf from tensorflow.core.protobuf import saved_model_pb2 from tensorflow.python.platform import gfile import easy_rec from easy_rec.python.builders import strategy_builder from easy_rec.python.compat import estimator_train from easy_rec.python.compat import exporter from easy_rec.python.input.input import Input from easy_rec.python.model.easy_rec_estimator import EasyRecEstimator from easy_rec.python.model.easy_rec_model import EasyRecModel from easy_rec.python.protos.train_pb2 import DistributionStrategy from easy_rec.python.utils import config_util from easy_rec.python.utils import constant from easy_rec.python.utils import estimator_utils from easy_rec.python.utils import fg_util from easy_rec.python.utils import load_class from easy_rec.python.utils.config_util import get_eval_input_path from easy_rec.python.utils.config_util import get_model_dir_path from easy_rec.python.utils.config_util import get_train_input_path from easy_rec.python.utils.config_util import set_eval_input_path from easy_rec.python.utils.export_big_model import export_big_model from easy_rec.python.utils.export_big_model import export_big_model_to_oss try: import horovod.tensorflow as hvd except Exception: hvd = None if tf.__version__ >= '2.0': from tensorflow.core.protobuf import config_pb2 ConfigProto = config_pb2.ConfigProto GPUOptions = config_pb2.GPUOptions tf = tf.compat.v1 else: GPUOptions = tf.GPUOptions ConfigProto = tf.ConfigProto load_class.auto_import() # when version of tensorflow > 1.8 strip_default_attrs set true will cause # saved_model inference core, such as: # [libprotobuf FATAL external/protobuf_archive/src/google/protobuf/map.h:1058] # CHECK failed: it != end(): key not found: new_axis_mask # so temporarily modify strip_default_attrs of _SavedModelExporter in # tf.estimator.exporter to false by default FinalExporter = exporter.FinalExporter LatestExporter = exporter.LatestExporter BestExporter = exporter.BestExporter def _get_input_fn(data_config, feature_configs, data_path=None, export_config=None, check_mode=False, **kwargs): """Build estimator input function. Args: data_config: dataset config feature_configs: FeatureConfig data_path: input_data_path export_config: configuration for exporting models, only used to build input_fn when exporting models Returns: subclass of Input """ input_class_map = {y: x for x, y in data_config.InputType.items()} input_cls_name = input_class_map[data_config.input_type] input_class = Input.create_class(input_cls_name) task_id, task_num = estimator_utils.get_task_index_and_num() input_obj = input_class( data_config, feature_configs, data_path, task_index=task_id, task_num=task_num, check_mode=check_mode, **kwargs) input_fn = input_obj.create_input(export_config) return input_fn def _create_estimator(pipeline_config, distribution=None, params={}): model_config = pipeline_config.model_config train_config = pipeline_config.train_config gpu_options = GPUOptions(allow_growth=True) # False) logging.info( 'train_config.train_distribute=%s[value=%d]' % (DistributionStrategy.Name(pipeline_config.train_config.train_distribute), pipeline_config.train_config.train_distribute)) # set gpu options only under hvd scenes if hvd is not None and pipeline_config.train_config.train_distribute in [ DistributionStrategy.EmbeddingParallelStrategy, DistributionStrategy.SokStrategy, DistributionStrategy.HorovodStrategy ]: local_rnk = hvd.local_rank() gpus = tf.config.experimental.list_physical_devices('GPU') logging.info('local_rnk=%d num_gpus=%d' % (local_rnk, len(gpus))) if len(gpus) > 0: tf.config.experimental.set_visible_devices(gpus[local_rnk], 'GPU') gpu_options.visible_device_list = str(local_rnk) session_config = ConfigProto( gpu_options=gpu_options, allow_soft_placement=True, log_device_placement=params.get('log_device_placement', False), inter_op_parallelism_threads=train_config.inter_op_parallelism_threads, intra_op_parallelism_threads=train_config.intra_op_parallelism_threads) if constant.NO_ARITHMETRIC_OPTI in os.environ: logging.info('arithmetic_optimization is closed to improve performance') session_config.graph_options.rewrite_options.arithmetic_optimization = \ session_config.graph_options.rewrite_options.OFF session_config.device_filters.append('/job:ps') model_cls = EasyRecModel.create_class(model_config.model_class) save_checkpoints_steps = None save_checkpoints_secs = None if train_config.HasField('save_checkpoints_steps'): save_checkpoints_steps = train_config.save_checkpoints_steps if train_config.HasField('save_checkpoints_secs'): save_checkpoints_secs = train_config.save_checkpoints_secs # if both `save_checkpoints_steps` and `save_checkpoints_secs` are not set, # use the default value of save_checkpoints_steps if save_checkpoints_steps is None and save_checkpoints_secs is None: save_checkpoints_steps = train_config.save_checkpoints_steps run_config = tf.estimator.RunConfig( model_dir=pipeline_config.model_dir, log_step_count_steps=None, # train_config.log_step_count_steps, save_summary_steps=train_config.save_summary_steps, save_checkpoints_steps=save_checkpoints_steps, save_checkpoints_secs=save_checkpoints_secs, keep_checkpoint_max=train_config.keep_checkpoint_max, train_distribute=distribution, eval_distribute=distribution, session_config=session_config) estimator = EasyRecEstimator( pipeline_config, model_cls, run_config=run_config, params=params) return estimator, run_config def _create_eval_export_spec(pipeline_config, eval_data, check_mode=False): data_config = pipeline_config.data_config # feature_configs = pipeline_config.feature_configs feature_configs = config_util.get_compatible_feature_configs(pipeline_config) eval_config = pipeline_config.eval_config export_config = pipeline_config.export_config if eval_config.num_examples > 0: eval_steps = int( math.ceil(float(eval_config.num_examples) / data_config.batch_size)) logging.info('eval_steps = %d' % eval_steps) else: eval_steps = None input_fn_kwargs = {'pipeline_config': pipeline_config} if data_config.input_type == data_config.InputType.OdpsRTPInputV2: input_fn_kwargs['fg_json_path'] = pipeline_config.fg_json_path # create eval input export_input_fn = _get_input_fn( data_config, feature_configs, None, export_config, check_mode=check_mode, **input_fn_kwargs) if export_config.exporter_type == 'final': exporters = [ FinalExporter(name='final', serving_input_receiver_fn=export_input_fn) ] elif export_config.exporter_type == 'latest': exporters = [ LatestExporter( name='latest', serving_input_receiver_fn=export_input_fn, exports_to_keep=export_config.exports_to_keep) ] elif export_config.exporter_type == 'best': logging.info( 'will use BestExporter, metric is %s, the bigger the better: %d' % (export_config.best_exporter_metric, export_config.metric_bigger)) def _metric_cmp_fn(best_eval_result, current_eval_result): logging.info('metric: best = %s current = %s' % (str(best_eval_result), str(current_eval_result))) if export_config.metric_bigger: return (best_eval_result[export_config.best_exporter_metric] < current_eval_result[export_config.best_exporter_metric]) else: return (best_eval_result[export_config.best_exporter_metric] > current_eval_result[export_config.best_exporter_metric]) exporters = [ BestExporter( name='best', serving_input_receiver_fn=export_input_fn, compare_fn=_metric_cmp_fn, exports_to_keep=export_config.exports_to_keep) ] elif export_config.exporter_type == 'none': exporters = [] else: raise ValueError('Unknown exporter type %s' % export_config.exporter_type) # set throttle_secs to a small number, so that we can control evaluation # interval steps by checkpoint saving steps eval_input_fn = _get_input_fn(data_config, feature_configs, eval_data, **input_fn_kwargs) eval_spec = tf.estimator.EvalSpec( name='val', input_fn=eval_input_fn, steps=eval_steps, throttle_secs=10, exporters=exporters) return eval_spec def _check_model_dir(model_dir, continue_train): if not continue_train: if not gfile.IsDirectory(model_dir): gfile.MakeDirs(model_dir) else: assert len(gfile.Glob(model_dir + '/model.ckpt-*.meta')) == 0, \ 'model_dir[=%s] already exists and not empty(if you ' \ 'want to continue train on current model_dir please ' \ 'delete dir %s or specify --continue_train[internal use only])' % ( model_dir, model_dir) else: if not gfile.IsDirectory(model_dir): logging.info('%s does not exists, create it automatically' % model_dir) gfile.MakeDirs(model_dir) def _get_ckpt_path(pipeline_config, checkpoint_path): if checkpoint_path != '' and checkpoint_path is not None: if gfile.IsDirectory(checkpoint_path): ckpt_path = estimator_utils.latest_checkpoint(checkpoint_path) else: ckpt_path = checkpoint_path elif gfile.IsDirectory(pipeline_config.model_dir): ckpt_path = estimator_utils.latest_checkpoint(pipeline_config.model_dir) logging.info('checkpoint_path is not specified, ' 'will use latest checkpoint %s from %s' % (ckpt_path, pipeline_config.model_dir)) else: assert False, 'pipeline_config.model_dir(%s) does not exist' \ % pipeline_config.model_dir return ckpt_path def train_and_evaluate(pipeline_config_path, continue_train=False): """Train and evaluate a EasyRec model defined in pipeline_config_path. Build an EasyRecEstimator, and then train and evaluate the estimator. Args: pipeline_config_path: a path to EasyRecConfig object, specifies train_config: model_config, data_config and eval_config continue_train: whether to restart train from an existing checkpoint Returns: None, the model will be saved into pipeline_config.model_dir """ assert gfile.Exists(pipeline_config_path), 'pipeline_config_path not exists' pipeline_config = config_util.get_configs_from_pipeline_file( pipeline_config_path) _train_and_evaluate_impl(pipeline_config, continue_train) return pipeline_config def _train_and_evaluate_impl(pipeline_config, continue_train=False, check_mode=False, fit_on_eval=False, fit_on_eval_steps=None): train_config = pipeline_config.train_config data_config = pipeline_config.data_config feature_configs = config_util.get_compatible_feature_configs(pipeline_config) if train_config.train_distribute != DistributionStrategy.NoStrategy\ and train_config.sync_replicas: logging.warning( 'will set sync_replicas to False, because train_distribute[%s] != NoStrategy' % pipeline_config.train_config.train_distribute) pipeline_config.train_config.sync_replicas = False train_data = get_train_input_path(pipeline_config) eval_data = get_eval_input_path(pipeline_config) distribution = strategy_builder.build(train_config) params = {} if train_config.is_profiling: params['log_device_placement'] = True estimator, run_config = _create_estimator( pipeline_config, distribution=distribution, params=params) version_file = os.path.join(pipeline_config.model_dir, 'version') if estimator_utils.is_chief(): _check_model_dir(pipeline_config.model_dir, continue_train) config_util.save_pipeline_config(pipeline_config, pipeline_config.model_dir) with gfile.GFile(version_file, 'w') as f: f.write(easy_rec.__version__ + '\n') train_steps = None if train_config.HasField('num_steps') and train_config.num_steps > 0: train_steps = train_config.num_steps assert train_steps is not None or data_config.num_epochs > 0, ( 'either num_steps and num_epochs must be set to an integer > 0.') if train_steps and data_config.num_epochs: logging.info('Both num_steps and num_epochs are set.') is_sync = train_config.sync_replicas batch_size = data_config.batch_size epoch_str = 'sample_num * %d / %d' % (data_config.num_epochs, batch_size) if is_sync: _, worker_num = estimator_utils.get_task_index_and_num() epoch_str += ' / ' + str(worker_num) logging.info('Will train min(%d, %s) steps...' % (train_steps, epoch_str)) input_fn_kwargs = {'pipeline_config': pipeline_config} if data_config.input_type == data_config.InputType.OdpsRTPInputV2: input_fn_kwargs['fg_json_path'] = pipeline_config.fg_json_path # create train input train_input_fn = _get_input_fn( data_config, feature_configs, train_data, check_mode=check_mode, **input_fn_kwargs) # Currently only a single Eval Spec is allowed. train_spec = tf.estimator.TrainSpec( input_fn=train_input_fn, max_steps=train_steps) embedding_parallel = train_config.train_distribute in ( DistributionStrategy.SokStrategy, DistributionStrategy.EmbeddingParallelStrategy) if embedding_parallel: estimator.train( input_fn=train_input_fn, max_steps=train_spec.max_steps, hooks=list(train_spec.hooks), saving_listeners=train_spec.saving_listeners) train_input_fn.input_creator.stop() else: # create eval spec eval_spec = _create_eval_export_spec( pipeline_config, eval_data, check_mode=check_mode) estimator_train.train_and_evaluate(estimator, train_spec, eval_spec) logging.info('Train and evaluate finish') if fit_on_eval and (not estimator_utils.is_evaluator()): tf.reset_default_graph() logging.info('Start continue training on eval data') eval_input_fn = _get_input_fn(data_config, feature_configs, eval_data, **input_fn_kwargs) if fit_on_eval_steps is not None: # wait estimator train done to get the correct train_steps while not estimator_train.estimator_train_done(estimator): time.sleep(1) train_steps = estimator_utils.get_trained_steps(estimator.model_dir) logging.info('\ttrain_steps=%d fit_on_eval_steps=%d' % (train_steps, fit_on_eval_steps)) fit_on_eval_steps += train_steps # Do not use estimator_train.train_and_evaluate as it starts tf.Server, # which is redundant and reports port not available error. estimator.train( input_fn=eval_input_fn, max_steps=fit_on_eval_steps, hooks=list(train_spec.hooks), saving_listeners=train_spec.saving_listeners if hasattr( train_spec, 'saving_listeners') else None) logging.info('Finished training on eval data') # return estimator for custom training using estimator.train return estimator def evaluate(pipeline_config, eval_checkpoint_path='', eval_data_path=None, eval_result_filename='eval_result.txt'): """Evaluate a EasyRec model defined in pipeline_config_path. Evaluate the model defined in pipeline_config_path on the eval data, the metrics will be displayed on tensorboard and saved into eval_result.txt. Args: pipeline_config: either EasyRecConfig path or its instance eval_checkpoint_path: if specified, will use this model instead of model specified by model_dir in pipeline_config_path eval_data_path: eval data path, default use eval data in pipeline_config could be a path or a list of paths eval_result_filename: evaluation result metrics save path. Returns: A dict of evaluation metrics: the metrics are specified in pipeline_config_path global_step: the global step for which this evaluation was performed. Raises: AssertionError, if: * pipeline_config_path does not exist """ pipeline_config = config_util.get_configs_from_pipeline_file(pipeline_config) if pipeline_config.fg_json_path: fg_util.load_fg_json_to_config(pipeline_config) if eval_data_path is not None: logging.info('Evaluating on data: %s' % eval_data_path) set_eval_input_path(pipeline_config, eval_data_path) train_config = pipeline_config.train_config eval_data = get_eval_input_path(pipeline_config) server_target = None if 'TF_CONFIG' in os.environ: tf_config = estimator_utils.chief_to_master() from tensorflow.python.training import server_lib if tf_config['task']['type'] == 'ps': cluster = tf.train.ClusterSpec(tf_config['cluster']) server = server_lib.Server( cluster, job_name='ps', task_index=tf_config['task']['index']) server.join() elif tf_config['task']['type'] == 'master': if 'ps' in tf_config['cluster']: cluster = tf.train.ClusterSpec(tf_config['cluster']) server = server_lib.Server(cluster, job_name='master', task_index=0) server_target = server.target print('server_target = %s' % server_target) distribution = strategy_builder.build(train_config) estimator, run_config = _create_estimator(pipeline_config, distribution) eval_spec = _create_eval_export_spec(pipeline_config, eval_data) ckpt_path = _get_ckpt_path(pipeline_config, eval_checkpoint_path) if server_target: # evaluate with parameter server input_iter = eval_spec.input_fn( mode=tf.estimator.ModeKeys.EVAL).make_one_shot_iterator() input_feas, input_lbls = input_iter.get_next() from tensorflow.python.training.device_setter import replica_device_setter from tensorflow.python.framework.ops import device from tensorflow.python.training.monitored_session import MonitoredSession from tensorflow.python.training.monitored_session import ChiefSessionCreator with device( replica_device_setter( worker_device='/job:master/task:0', cluster=cluster)): estimator_spec = estimator._eval_model_fn(input_feas, input_lbls, run_config) session_config = ConfigProto( allow_soft_placement=True, log_device_placement=True) chief_sess_creator = ChiefSessionCreator( master=server_target, checkpoint_filename_with_path=ckpt_path, config=session_config) eval_metric_ops = estimator_spec.eval_metric_ops update_ops = [eval_metric_ops[x][1] for x in eval_metric_ops.keys()] metric_ops = {x: eval_metric_ops[x][0] for x in eval_metric_ops.keys()} update_op = tf.group(update_ops) with MonitoredSession( session_creator=chief_sess_creator, hooks=None, stop_grace_period_secs=120) as sess: while True: try: sess.run(update_op) except tf.errors.OutOfRangeError: break eval_result = sess.run(metric_ops) else: # this way does not work, wait to be debugged # the variables are not placed to parameter server # with tf.device( # replica_device_setter( # worker_device='/job:master/task:0', cluster=cluster)): eval_result = estimator.evaluate( eval_spec.input_fn, eval_spec.steps, checkpoint_path=ckpt_path) eval_spec.input_fn.input_creator.stop() logging.info('Evaluate finish') print('eval_result = ', eval_result) logging.info('eval_result = {0}'.format(eval_result)) # write eval result to file model_dir = pipeline_config.model_dir eval_result_file = os.path.join(model_dir, eval_result_filename) logging.info('save eval result to file %s' % eval_result_file) with gfile.GFile(eval_result_file, 'w') as ofile: result_to_write = {} for key in sorted(eval_result): # skip logging binary data if isinstance(eval_result[key], six.binary_type): continue # convert numpy float to python float result_to_write[key] = eval_result[key].item() ofile.write(json.dumps(result_to_write, indent=2)) return eval_result def distribute_evaluate(pipeline_config, eval_checkpoint_path='', eval_data_path=None, eval_result_filename='distribute_eval_result.txt'): """Evaluate a EasyRec model defined in pipeline_config_path. Evaluate the model defined in pipeline_config_path on the eval data, the metrics will be displayed on tensorboard and saved into eval_result.txt. Args: pipeline_config: either EasyRecConfig path or its instance eval_checkpoint_path: if specified, will use this model instead of model specified by model_dir in pipeline_config_path eval_data_path: eval data path, default use eval data in pipeline_config could be a path or a list of paths eval_result_filename: evaluation result metrics save path. Returns: A dict of evaluation metrics: the metrics are specified in pipeline_config_path global_step: the global step for which this evaluation was performed. Raises: AssertionError, if: * pipeline_config_path does not exist """ pipeline_config = config_util.get_configs_from_pipeline_file(pipeline_config) if eval_data_path is not None: logging.info('Evaluating on data: %s' % eval_data_path) set_eval_input_path(pipeline_config, eval_data_path) train_config = pipeline_config.train_config eval_data = get_eval_input_path(pipeline_config) data_config = pipeline_config.data_config if data_config.HasField('sampler'): logging.warning( 'It is not accuracy to use eval with negative sampler, recommand to use hitrate.py!' ) eval_result = {} return eval_result model_dir = get_model_dir_path(pipeline_config) eval_tmp_results_dir = os.path.join(model_dir, 'distribute_eval_tmp_results') if not gfile.IsDirectory(eval_tmp_results_dir): logging.info('create eval tmp results dir {}'.format(eval_tmp_results_dir)) gfile.MakeDirs(eval_tmp_results_dir) assert gfile.IsDirectory( eval_tmp_results_dir), 'tmp results dir not create success.' os.environ['eval_tmp_results_dir'] = eval_tmp_results_dir server_target = None cur_job_name = None if 'TF_CONFIG' in os.environ: tf_config = estimator_utils.chief_to_master() from tensorflow.python.training import server_lib if tf_config['task']['type'] == 'ps': cluster = tf.train.ClusterSpec(tf_config['cluster']) server = server_lib.Server( cluster, job_name='ps', task_index=tf_config['task']['index']) server.join() elif tf_config['task']['type'] == 'master': if 'ps' in tf_config['cluster']: cur_job_name = tf_config['task']['type'] cur_task_index = tf_config['task']['index'] cluster = tf.train.ClusterSpec(tf_config['cluster']) server = server_lib.Server( cluster, job_name=cur_job_name, task_index=cur_task_index) server_target = server.target print('server_target = %s' % server_target) elif tf_config['task']['type'] == 'worker': if 'ps' in tf_config['cluster']: cur_job_name = tf_config['task']['type'] cur_task_index = tf_config['task']['index'] cluster = tf.train.ClusterSpec(tf_config['cluster']) server = server_lib.Server( cluster, job_name=cur_job_name, task_index=cur_task_index) server_target = server.target print('server_target = %s' % server_target) if server_target: from tensorflow.python.training.device_setter import replica_device_setter from tensorflow.python.framework.ops import device from tensorflow.python.training.monitored_session import MonitoredSession from tensorflow.python.training.monitored_session import ChiefSessionCreator from tensorflow.python.training.monitored_session import WorkerSessionCreator from easy_rec.python.utils.estimator_utils import EvaluateExitBarrierHook cur_work_device = '/job:' + cur_job_name + '/task:' + str(cur_task_index) cur_ps_num = len(tf_config['cluster']['ps']) with device( replica_device_setter( ps_tasks=cur_ps_num, worker_device=cur_work_device, cluster=cluster)): distribution = strategy_builder.build(train_config) estimator, run_config = _create_estimator(pipeline_config, distribution) eval_spec = _create_eval_export_spec(pipeline_config, eval_data) ckpt_path = _get_ckpt_path(pipeline_config, eval_checkpoint_path) ckpt_dir = os.path.dirname(ckpt_path) input_iter = eval_spec.input_fn( mode=tf.estimator.ModeKeys.EVAL).make_one_shot_iterator() input_feas, input_lbls = input_iter.get_next() estimator_spec = estimator._distribute_eval_model_fn( input_feas, input_lbls, run_config) session_config = ConfigProto( allow_soft_placement=True, log_device_placement=True, device_filters=['/job:ps', '/job:worker/task:%d' % cur_task_index]) if cur_job_name == 'master': metric_variables = tf.get_collection(tf.GraphKeys.METRIC_VARIABLES) model_ready_for_local_init_op = tf.variables_initializer(metric_variables) global_variables = tf.global_variables() remain_variables = list( set(global_variables).difference(set(metric_variables))) cur_saver = tf.train.Saver(var_list=remain_variables, sharded=True) cur_scaffold = tf.train.Scaffold( saver=cur_saver, ready_for_local_init_op=model_ready_for_local_init_op) cur_sess_creator = ChiefSessionCreator( scaffold=cur_scaffold, master=server_target, checkpoint_filename_with_path=ckpt_path, config=session_config) else: cur_sess_creator = WorkerSessionCreator( master=server_target, config=session_config) eval_metric_ops = estimator_spec.eval_metric_ops update_ops = [eval_metric_ops[x][1] for x in eval_metric_ops.keys()] metric_ops = {x: eval_metric_ops[x][0] for x in eval_metric_ops.keys()} update_op = tf.group(update_ops) cur_worker_num = len(tf_config['cluster']['worker']) + 1 if cur_job_name == 'master': cur_stop_grace_period_sesc = 120 cur_hooks = EvaluateExitBarrierHook(cur_worker_num, True, ckpt_dir, metric_ops) else: cur_stop_grace_period_sesc = 10 cur_hooks = EvaluateExitBarrierHook(cur_worker_num, False, ckpt_dir, metric_ops) with MonitoredSession( session_creator=cur_sess_creator, hooks=[cur_hooks], stop_grace_period_secs=cur_stop_grace_period_sesc) as sess: while True: try: sess.run(update_op) except tf.errors.OutOfRangeError: break eval_result = cur_hooks.eval_result logging.info('Evaluate finish') # write eval result to file model_dir = pipeline_config.model_dir eval_result_file = os.path.join(model_dir, eval_result_filename) logging.info('save eval result to file %s' % eval_result_file) if cur_job_name == 'master': print('eval_result = ', eval_result) logging.info('eval_result = {0}'.format(eval_result)) with gfile.GFile(eval_result_file, 'w') as ofile: result_to_write = {'eval_method': 'distribute'} for key in sorted(eval_result): # skip logging binary data if isinstance(eval_result[key], six.binary_type): continue # convert numpy float to python float result_to_write[key] = eval_result[key].item() ofile.write(json.dumps(result_to_write)) return eval_result def predict(pipeline_config, checkpoint_path='', data_path=None): """Predict a EasyRec model defined in pipeline_config_path. Predict the model defined in pipeline_config_path on the eval data. Args: pipeline_config: either EasyRecConfig path or its instance checkpoint_path: if specified, will use this model instead of model specified by model_dir in pipeline_config_path data_path: data path, default use eval data in pipeline_config could be a path or a list of paths Returns: A list of dict of predict results Raises: AssertionError, if: * pipeline_config_path does not exist """ pipeline_config = config_util.get_configs_from_pipeline_file(pipeline_config) if pipeline_config.fg_json_path: fg_util.load_fg_json_to_config(pipeline_config) if data_path is not None: logging.info('Predict on data: %s' % data_path) set_eval_input_path(pipeline_config, data_path) train_config = pipeline_config.train_config eval_data = get_eval_input_path(pipeline_config) distribution = strategy_builder.build(train_config) estimator, _ = _create_estimator(pipeline_config, distribution) eval_spec = _create_eval_export_spec(pipeline_config, eval_data) ckpt_path = _get_ckpt_path(pipeline_config, checkpoint_path) pred_result = estimator.predict(eval_spec.input_fn, checkpoint_path=ckpt_path) logging.info('Predict finish') return pred_result def export(export_dir, pipeline_config, checkpoint_path='', asset_files=None, verbose=False, **extra_params): """Export model defined in pipeline_config_path. Args: export_dir: base directory where the model should be exported pipeline_config: proto.EasyRecConfig instance or file path specify proto.EasyRecConfig checkpoint_path: if specified, will use this model instead of model in model_dir in pipeline_config_path asset_files: extra files to add to assets, comma separated; if asset file variable in graph need to be renamed, specify by new_file_name:file_path version: if version is defined, then will skip writing embedding to redis, assume that embedding is already write into redis verbose: dumps debug information extra_params: keys related to write embedding to redis/oss redis_url, redis_passwd, redis_threads, redis_batch_size, redis_timeout, redis_expire if export embedding to redis; redis_embedding_version: if specified, will kill export to redis -- oss_path, oss_endpoint, oss_ak, oss_sk, oss_timeout, oss_expire, oss_write_kv, oss_embedding_version Returns: the directory where model is exported Raises: AssertionError, if: * pipeline_config_path does not exist """ if not gfile.Exists(export_dir): gfile.MakeDirs(export_dir) pipeline_config = config_util.get_configs_from_pipeline_file(pipeline_config) if pipeline_config.fg_json_path: fg_util.load_fg_json_to_config(pipeline_config) # feature_configs = pipeline_config.feature_configs feature_configs = config_util.get_compatible_feature_configs(pipeline_config) # create estimator params = {'log_device_placement': verbose} if asset_files: logging.info('will add asset files: %s' % asset_files) asset_file_dict = {} for asset_file in asset_files.split(','): asset_file = asset_file.strip() if ':' not in asset_file or asset_file.startswith( 'oss:') or asset_file.startswith('hdfs:'): _, asset_name = os.path.split(asset_file) else: asset_name, asset_file = asset_file.split(':', 1) asset_file_dict[asset_name] = asset_file params['asset_files'] = asset_file_dict estimator, _ = _create_estimator(pipeline_config, params=params) # construct serving input fn export_config = pipeline_config.export_config data_config = pipeline_config.data_config input_fn_kwargs = {'pipeline_config': pipeline_config} if data_config.input_type == data_config.InputType.OdpsRTPInputV2: input_fn_kwargs['fg_json_path'] = pipeline_config.fg_json_path serving_input_fn = _get_input_fn(data_config, feature_configs, None, export_config, **input_fn_kwargs) ckpt_path = _get_ckpt_path(pipeline_config, checkpoint_path) if 'oss_path' in extra_params: if pipeline_config.train_config.HasField('incr_save_config'): incr_save_config = pipeline_config.train_config.incr_save_config extra_params['incr_update'] = {} incr_save_type = incr_save_config.WhichOneof('incr_update') logging.info('incr_save_type=%s' % incr_save_type) if incr_save_type: extra_params['incr_update'][incr_save_type] = getattr( incr_save_config, incr_save_type) return export_big_model_to_oss(export_dir, pipeline_config, extra_params, serving_input_fn, estimator, ckpt_path, verbose) if 'redis_url' in extra_params: return export_big_model(export_dir, pipeline_config, extra_params, serving_input_fn, estimator, ckpt_path, verbose) final_export_dir = estimator.export_savedmodel( export_dir_base=export_dir, serving_input_receiver_fn=serving_input_fn, checkpoint_path=ckpt_path, strip_default_attrs=True) # add export ts as version info saved_model = saved_model_pb2.SavedModel() if type(final_export_dir) not in [type(''), type(u'')]: final_export_dir = final_export_dir.decode('utf-8') export_ts = [ x for x in final_export_dir.split('/') if x != '' and x is not None ] export_ts = export_ts[-1] saved_pb_path = os.path.join(final_export_dir, 'saved_model.pb') with gfile.GFile(saved_pb_path, 'rb') as fin: saved_model.ParseFromString(fin.read()) saved_model.meta_graphs[0].meta_info_def.meta_graph_version = export_ts with gfile.GFile(saved_pb_path, 'wb') as fout: fout.write(saved_model.SerializeToString()) logging.info('model has been exported to %s successfully' % final_export_dir) return final_export_dir def export_checkpoint(pipeline_config=None, export_path='', checkpoint_path='', asset_files=None, verbose=False, mode=tf.estimator.ModeKeys.PREDICT): """Export the EasyRec model as checkpoint.""" pipeline_config = config_util.get_configs_from_pipeline_file(pipeline_config) if pipeline_config.fg_json_path: fg_util.load_fg_json_to_config(pipeline_config) feature_configs = config_util.get_compatible_feature_configs(pipeline_config) data_config = pipeline_config.data_config input_fn_kwargs = {'pipeline_config': pipeline_config} if data_config.input_type == data_config.InputType.OdpsRTPInputV2: input_fn_kwargs['fg_json_path'] = pipeline_config.fg_json_path # create estimator params = {'log_device_placement': verbose} if asset_files: logging.info('will add asset files: %s' % asset_files) params['asset_files'] = asset_files estimator, _ = _create_estimator(pipeline_config, params=params) # construct serving input fn export_config = pipeline_config.export_config serving_input_fn = _get_input_fn(data_config, feature_configs, None, export_config, **input_fn_kwargs) ckpt_path = _get_ckpt_path(pipeline_config, checkpoint_path) estimator.export_checkpoint( export_path=export_path, serving_input_receiver_fn=serving_input_fn, checkpoint_path=ckpt_path, mode=mode) logging.info('model checkpoint has been exported successfully')