def testGetFnFailure(self): exec_properties = { 'module_file': 'path/to/module_file.py', 'test_fn': 'path.to.test_fn', } with self.assertRaises(ValueError): udf_utils.get_fn(exec_properties, 'test_fn')
def Do(self, input_dict: Dict[Text, List[types.Artifact]], output_dict: Dict[Text, List[types.Artifact]], exec_properties: Dict[Text, Any]) -> None: if exec_properties.get(_TUNE_ARGS_KEY): raise ValueError( "TuneArgs is not supported for default Tuner's Executor.") tuner_fn = udf_utils.get_fn(exec_properties, 'tuner_fn') fn_args = fn_args_utils.get_common_fn_args(input_dict, exec_properties, self._get_tmp_dir()) tuner_fn_result = tuner_fn(fn_args) tuner = tuner_fn_result.tuner fit_kwargs = tuner_fn_result.fit_kwargs # TODO(b/156966497): set logger for printing. tuner.search_space_summary() absl.logging.info('Start tuning...') tuner.search(**fit_kwargs) tuner.results_summary() best_hparams_config = tuner.get_best_hyperparameters()[0].get_config() absl.logging.info('Best hyperParameters: %s' % best_hparams_config) best_hparams_path = os.path.join( artifact_utils.get_single_uri( output_dict[_BEST_HYPERPARAMETERS_KEY]), _DEFAULT_FILE_NAME) io_utils.write_string_file(best_hparams_path, json.dumps(best_hparams_config)) absl.logging.info('Best Hyperparameters are written to %s.' % best_hparams_path)
def Do(self, input_dict: Dict[str, List[types.Artifact]], output_dict: Dict[str, List[types.Artifact]], exec_properties: Dict[str, Any]) -> None: if tfx_tuner.get_tune_args(exec_properties): raise ValueError( "TuneArgs is not supported by this Tuner's Executor.") metalearning_algorithm = None if 'metalearning_algorithm' in exec_properties: metalearning_algorithm = exec_properties.get( 'metalearning_algorithm') warmup_trials = 0 warmup_trial_data = None if metalearning_algorithm: warmup_tuner, warmup_trials = self.warmup(input_dict, exec_properties, metalearning_algorithm) warmup_trial_data = extract_tuner_trial_progress(warmup_tuner) else: logging.info('MetaLearning Algorithm not provided.') # Create new fn_args for final tuning stage. fn_args = fn_args_utils.get_common_fn_args( input_dict, exec_properties, working_dir=self._get_tmp_dir()) tuner_fn = udf_utils.get_fn(exec_properties, 'tuner_fn') tuner_fn_result = tuner_fn(fn_args) tuner_fn_result.tuner.oracle.max_trials = max( (tuner_fn_result.tuner.oracle.max_trials - warmup_trials), 1) tuner = self.search(tuner_fn_result) tuner_trial_data = extract_tuner_trial_progress(tuner) if warmup_trial_data: cumulative_tuner_trial_data, best_tuner_ix = merge_trial_data( warmup_trial_data, tuner_trial_data) cumulative_tuner_trial_data[ 'warmup_trial_data'] = warmup_trial_data[BEST_CUMULATIVE_SCORE] cumulative_tuner_trial_data['tuner_trial_data'] = tuner_trial_data[ BEST_CUMULATIVE_SCORE] if isinstance(tuner.oracle.objective, kerastuner.Objective): cumulative_tuner_trial_data[ 'objective'] = tuner.oracle.objective.name else: cumulative_tuner_trial_data[ 'objective'] = 'objective not understood' tuner_trial_data = cumulative_tuner_trial_data best_tuner = warmup_tuner if best_tuner_ix == 0 else tuner else: best_tuner = tuner tfx_tuner.write_best_hyperparameters(best_tuner, output_dict) tuner_plot_path = os.path.join( artifact_utils.get_single_uri(output_dict['trial_summary_plot']), 'tuner_plot_data.txt') io_utils.write_string_file(tuner_plot_path, json.dumps(tuner_trial_data)) logging.info('Tuner plot data written at: %s', tuner_plot_path)
def Do(self, input_dict: Dict[Text, List[types.Artifact]], output_dict: Dict[Text, List[types.Artifact]], exec_properties: Dict[Text, Any]) -> None: self._log_startup(input_dict, output_dict, exec_properties) create_decoder_func = udf_utils.get_fn(exec_properties, _CREATE_DECODER_FUNC_KEY) tf_graph_record_decoder.save_decoder( create_decoder_func(), value_utils.GetSoleValue(output_dict, _DATA_VIEW_KEY).uri)
def Do(self, input_dict: Dict[Text, List[types.Artifact]], output_dict: Dict[Text, List[types.Artifact]], exec_properties: Dict[Text, Any]) -> None: """Uses a user-supplied run_fn to train a TensorFlow model locally. The Trainer Executor invokes a run_fn callback function provided by the user via the module_file parameter. In this function, user defines the model and trains it, then saves the model and training related files (e.g, Tensorboard logs) to the provided locations. Args: input_dict: Input dict from input key to a list of ML-Metadata Artifacts. - examples: Examples used for training, must include 'train' and 'eval' if custom splits is not specified in train_args and eval_args. - transform_graph: Optional input transform graph. - transform_output: Optional input transform graph, deprecated. - schema: Schema of the data. output_dict: Output dict from output key to a list of Artifacts. - model: Exported model. - model_run: Model training related outputs (e.g., Tensorboard logs) exec_properties: A dict of execution properties. - train_args: JSON string of trainer_pb2.TrainArgs instance, providing args for training. - eval_args: JSON string of trainer_pb2.EvalArgs instance, providing args for eval. - module_file: Python module file containing UDF model definition. - warm_starting: Whether or not we need to do warm starting. - warm_start_from: Optional. If warm_starting is True, this is the directory to find previous model to warm start on. - custom_config: Optional. JSON-serialized dict of additional parameters to pass to trainer function. Returns: None Raises: ValueError: When neither or both of 'module_file' and 'run_fn' are present in 'exec_properties'. RuntimeError: If run_fn failed to generate model in desired location. """ self._log_startup(input_dict, output_dict, exec_properties) fn_args = self._GetFnArgs(input_dict, output_dict, exec_properties) run_fn = udf_utils.get_fn(exec_properties, 'run_fn') # Train the model absl.logging.info('Training model.') run_fn(fn_args) # Note: If trained with multi-node distribution workers, it is the user # module's responsibility to export the model only once. if not fileio.exists(fn_args.serving_model_dir): raise RuntimeError('run_fn failed to generate model.') absl.logging.info( 'Training complete. Model written to %s. ModelRun written to %s', fn_args.serving_model_dir, fn_args.model_run_dir)
def Do(self, input_dict: Dict[Text, List[types.Artifact]], output_dict: Dict[Text, List[types.Artifact]], exec_properties: Dict[Text, Any]) -> None: del input_dict if _MODULE_FILE_KEY in exec_properties: create_decoder_func = import_utils.import_func_from_source( exec_properties.get(_MODULE_FILE_KEY), exec_properties.get(_CREATE_DECODER_FUNC_KEY)) else: create_decoder_func = udf_utils.get_fn(exec_properties, _CREATE_DECODER_FUNC_KEY) tf_graph_record_decoder.save_decoder( create_decoder_func(), value_utils.GetSoleValue(output_dict, _DATA_VIEW_KEY).uri)
def _run_transform(args, beam_pipeline_args): """Construct and run transform executor.""" absl.logging.set_verbosity(absl.logging.INFO) def make_beam_pipeline(): return beam.Pipeline(beam_pipeline_args) preprocessing_fn = udf_utils.get_fn( { standard_component_specs.PREPROCESSING_FN_KEY: args.preprocessing_fn_path }, standard_component_specs.PREPROCESSING_FN_KEY) inputs = { labels.ANALYZE_DATA_PATHS_LABEL: args.analyze_examples, labels.ANALYZE_PATHS_FILE_FORMATS_LABEL: [labels.FORMAT_TFRECORD] * len(args.analyze_examples), labels.TRANSFORM_DATA_PATHS_LABEL: [args.analyze_examples + args.transform_only_examples], labels.TRANSFORM_PATHS_FILE_FORMATS_LABEL: [labels.FORMAT_TFRECORD] * (len(args.analyze_examples) + len(args.transform_only_examples)), labels.SCHEMA_PATH_LABEL: args.input_schema_path, labels.PREPROCESSING_FN: preprocessing_fn, labels.EXAMPLES_DATA_FORMAT_LABEL: example_gen_pb2.PayloadFormat.Value(args.example_data_format), labels.DISABLE_STATISTICS_LABEL: args.disable_statistics, labels.MAKE_BEAM_PIPELINE_FN: make_beam_pipeline, } outputs = { labels.TRANSFORM_METADATA_OUTPUT_PATH_LABEL: args.transform_fn, labels.TRANSFORM_MATERIALIZE_OUTPUT_PATHS_LABEL: (args.transformed_examples), labels.PER_SET_STATS_OUTPUT_PATHS_LABEL: (args.per_set_stats_outputs), labels.TEMP_OUTPUT_LABEL: args.tmp_location, } executor.TransformProcessor().Transform(inputs, outputs, args.status_file)
def testGetFnFromModule(self, mock_import_func): exec_properties = {'test_fn': 'path.to.test_fn'} udf_utils.get_fn(exec_properties, 'test_fn') mock_import_func.assert_called_once_with('path.to', 'test_fn')
def testGetFnFromSource(self, mock_import_func): exec_properties = {'module_file': 'path/to/module_file.py'} udf_utils.get_fn(exec_properties, 'test_fn') mock_import_func.assert_called_once_with('path/to/module_file.py', 'test_fn')
def testGetFnFailure(self): with self.assertRaises(ValueError): udf_utils.get_fn({}, 'test_fn')
def Do(self, input_dict: Dict[Text, List[types.Artifact]], output_dict: Dict[Text, List[types.Artifact]], exec_properties: Dict[Text, Any]) -> None: """Uses a user-supplied tf.estimator to train a TensorFlow model locally. The Trainer Executor invokes a training_fn callback function provided by the user via the module_file parameter. With the tf.estimator returned by this function, the Trainer Executor then builds a TensorFlow model using the user-provided tf.estimator. Args: input_dict: Input dict from input key to a list of ML-Metadata Artifacts. - examples: Examples used for training, must include 'train' and 'eval' splits. - transform_output: Optional input transform graph. - schema: Schema of the data. output_dict: Output dict from output key to a list of Artifacts. - model: Exported model. - model_run: Model training related outputs (e.g., Tensorboard logs) exec_properties: A dict of execution properties. - train_args: JSON string of trainer_pb2.TrainArgs instance, providing args for training. - eval_args: JSON string of trainer_pb2.EvalArgs instance, providing args for eval. - module_file: Python module file containing UDF model definition. - warm_starting: Whether or not we need to do warm starting. - warm_start_from: Optional. If warm_starting is True, this is the directory to find previous model to warm start on. - custom_config: Optional. JSON-serialized dict of additional parameters to pass to trainer function. Returns: None Raises: ValueError: When neither or both of 'module_file' and 'trainer_fn' are present in 'exec_properties'. """ self._log_startup(input_dict, output_dict, exec_properties) fn_args = self._GetFnArgs(input_dict, output_dict, exec_properties) trainer_fn = udf_utils.get_fn(exec_properties, 'trainer_fn') schema = io_utils.parse_pbtxt_file(fn_args.schema_file, schema_pb2.Schema()) # TODO(b/160795287): Deprecate estimator based executor. # Provide user with a modified fn_args, with model_run given as # the working directory. Executor will then copy user models to # model artifact directory. serving_dest = fn_args.serving_model_dir eval_dest = fn_args.eval_model_dir working_dir = fn_args.model_run_dir fn_args.serving_model_dir = path_utils.serving_model_dir(working_dir) fn_args.eval_model_dir = path_utils.eval_model_dir(working_dir) training_spec = trainer_fn(fn_args, schema) # Train the model absl.logging.info('Training model.') tf.estimator.train_and_evaluate(training_spec['estimator'], training_spec['train_spec'], training_spec['eval_spec']) absl.logging.info( 'Training complete. Model written to %s. ModelRun written to %s', fn_args.serving_model_dir, fn_args.model_run_dir) # Export an eval savedmodel for TFMA. If distributed training, it must only # be written by the chief worker, as would be done for serving savedmodel. if _is_chief(): absl.logging.info('Exporting eval_savedmodel for TFMA.') tfma.export.export_eval_savedmodel( estimator=training_spec['estimator'], export_dir_base=fn_args.eval_model_dir, eval_input_receiver_fn=training_spec['eval_input_receiver_fn']) absl.logging.info('Exported eval_savedmodel to %s.', fn_args.eval_model_dir) # TODO(b/160795287): Deprecate estimator based executor. # Copy serving and eval model from model_run to model artifact directory. serving_source = path_utils.serving_model_path(fn_args.model_run_dir) io_utils.copy_dir(serving_source, serving_dest) absl.logging.info('Serving model copied to: %s.', serving_dest) eval_source = path_utils.eval_model_path(fn_args.model_run_dir) io_utils.copy_dir(eval_source, eval_dest) absl.logging.info('Eval model copied to: %s.', eval_dest) else: absl.logging.info( 'Model export is skipped because this is not the chief worker.')
def warmup(self, input_dict: Dict[str, List[types.Artifact]], exec_properties: Dict[str, List[types.Artifact]], algorithm: str): # Perform warmup tuning if WARMUP_HYPERPARAMETERS given. hparams_warmup_config_list = None if input_dict.get(WARMUP_HYPERPARAMETERS): hyperparameters_file = io_utils.get_only_uri_in_dir( artifact_utils.get_single_uri( input_dict[WARMUP_HYPERPARAMETERS])) hparams_warmup_config_list = json.loads( io_utils.read_string_file(hyperparameters_file)) fn_args = fn_args_utils.get_common_fn_args( input_dict, exec_properties, working_dir=self._get_tmp_dir() + 'warmup') # TODO(nikhilmehta): Currently all algorithms need warmup_hyperparameters. # This may not be needed for other algorithms that can predict hyperparams. if not hparams_warmup_config_list: raise ValueError('Expected warmup_hyperparameters') logging.info('Algorithm: %s', algorithm) warmup_trials = 0 if algorithm == 'majority_voting': warmup_trials = DEFAULT_WARMUP_TRIALS fn_args.custom_config[ WARMUP_HYPERPARAMETERS] = hparams_warmup_config_list[0] elif algorithm == 'nearest_neighbor': warmup_trials = DEFAULT_WARMUP_TRIALS if input_dict.get('metamodel'): metamodel_path = io_utils.get_only_uri_in_dir( artifact_utils.get_single_uri(input_dict['metamodel'])) logging.info('Meta model path: %s', metamodel_path) metamodel = _load_keras_model(metamodel_path) else: raise ValueError( f'Tuner for metalearning_algorithm={algorithm} expects metamodel.' ) if input_dict.get('metafeature'): metafeature_path = io_utils.get_only_uri_in_dir( artifact_utils.get_single_uri(input_dict['metafeature'])) logging.info('Metafeature: %s', metafeature_path) metafeature = json.loads( io_utils.read_string_file(metafeature_path)) metafeature = metafeature['metafeature'] else: raise ValueError( f'Tuner for metalearning_algorithm={algorithm} expects metafeature.' ) metafeature = np.array(metafeature, dtype=np.float32) metafeature = np.expand_dims(metafeature, axis=0) logits = metamodel(metafeature).numpy()[0] nearest_configs = [ hparams_warmup_config_list[ix] for ix in np.argsort(logits)[-DEFAULT_K:] ] nearest_hparam_config = _merge_hparam_configs(nearest_configs) fn_args.custom_config[ WARMUP_HYPERPARAMETERS] = nearest_hparam_config else: raise NotImplementedError( f'Tuning for metalearning_algorithm={algorithm} is not implemented.' ) # kerastuner doesn't support grid search, setting max_trials large enough. # Track issue: https://github.com/keras-team/keras-tuner/issues/340 fn_args.custom_config['max_trials'] = warmup_trials tuner_fn = udf_utils.get_fn(exec_properties, 'tuner_fn') warmtuner_fn_result = tuner_fn(fn_args) warmup_tuner = self.search(warmtuner_fn_result) return warmup_tuner, warmup_trials
def _get_tuner_fn(exec_properties: Dict[str, Any]) -> Callable[..., Any]: """Returns tuner_fn from execution properties.""" return udf_utils.get_fn(exec_properties, 'tuner_fn')