def convert(assets_filepath): """Converts existing asset pickle based files to t2r proto based assets.""" t2r_assets = t2r_pb2.T2RAssets() input_spec_filepath = os.path.join(assets_filepath, 'input_specs.pkl') if not tf.io.gfile.exists(input_spec_filepath): raise ValueError('No file exists for {}.'.format(input_spec_filepath)) feature_spec, label_spec = tensorspec_utils.load_input_spec_from_file( input_spec_filepath) t2r_assets.feature_spec.CopyFrom(feature_spec.to_proto()) t2r_assets.label_spec.CopyFrom(label_spec.to_proto()) global_step_filepath = os.path.join(assets_filepath, 'global_step.pkl') if tf.io.gfile.exists(global_step_filepath): global_step = tensorspec_utils.load_input_spec_from_file( global_step_filepath) t2r_assets.global_step = global_step t2r_assets_filepath = os.path.join(assets_filepath, tensorspec_utils.T2R_ASSETS_FILENAME) tensorspec_utils.write_t2r_assets_to_file(t2r_assets, t2r_assets_filepath)
def restore(self): """Restores the model parameters from the latest available data. Raises: ValueError: If no checkpoint can be found or loaded within the user defined timeout. Returns: True if a exported saved model has been loaded and False otherwise. """ start_time = time.time() while time.time() - start_time < self._timeout: # The exported saved models directory names are numbers (timestamp) which # monotonically increase, meaning the largest directory name will contain # the latest exported model. Lexicographical sorting will maintain this # order. model_dirs = sorted(tf.io.gfile.glob(os.path.join(self._export_dir, '*'))) model_dirs = self._remove_invalid_model_dirnames(model_dirs) if len(model_dirs) >= 1: logging.info('Found latest model at %s. ', model_dirs[-1]) break logging.info('Waiting for an exported model to become available at %s.', self._export_dir) # Since a checkpoint might not be available and this is a busy waiting # loop, we throttle checking for checkpoints. time.sleep(_BUSY_WAITING_SLEEP_TIME_IN_SECS) if model_dirs is None or not model_dirs: logging.warning('No checkpoint available after %s seconds.', str(self._timeout)) return False if self._latest_export_dir == model_dirs[-1]: # The latest model has already been loaded. return True logging.info('Loading the latest model at %s. ', model_dirs[-1]) self._latest_export_dir = model_dirs[-1] start_time_loading = time.time() # Note, loading from a saved model might require several attempts if # the checkpoint gets written asynchronously. while time.time() - start_time_loading < self._timeout: try: self._predict_fn = tf.contrib.predictor.from_saved_model( model_dirs[-1], config=self._tf_config) t2r_assets_file_path = os.path.join( model_dirs[-1], 'assets.extra', tensorspec_utils.T2R_ASSETS_FILENAME) if tf.io.gfile.exists(t2r_assets_file_path): t2r_assets = tensorspec_utils.load_t2r_assets_to_file( t2r_assets_file_path) self._feature_spec = tensorspec_utils.TensorSpecStruct.from_proto( t2r_assets.feature_spec) # pytype: disable=wrong-arg-types self._label_spec = tensorspec_utils.TensorSpecStruct.from_proto( t2r_assets.label_spec) # pytype: disable=wrong-arg-types if t2r_assets.HasField('global_step'): self._global_step = t2r_assets.global_step else: logging.warning( 'Error loading the global step, therefore using the previously' 'set global step %s.', str(self.global_step)) else: input_spec_filename = os.path.join(model_dirs[-1], 'assets.extra', 'input_specs.pkl') logging.warning( 'Using the legacy loading, please convert the assets ' 'using convert_pkl_assets_to_proto_assets binary for ' 'file path %s.', input_spec_filename) # Load input specs from file. self._feature_spec, self._label_spec = ( tensorspec_utils.load_input_spec_from_file(input_spec_filename)) # Load input specs from file. global_step_filename = os.path.join(model_dirs[-1], 'assets.extra', 'global_step.pkl') try: global_step = tensorspec_utils.load_global_step_from_file( global_step_filename) self._global_step = global_step except ValueError: logging.warning( 'Error loading the global step, therefore using the previously' 'set global step %s.', str(self.global_step)) return True except ValueError as err: logging.warning( 'Error loading model as %s:\n%s\nThe next attempt at loading the ' 'latest model will be in %d seconds', model_dirs[-1], err, _BUSY_WAITING_SLEEP_TIME_IN_SECS) # Since a checkpoint might be written by the tf model concurrently # this is a busy waiting loop. time.sleep(_BUSY_WAITING_SLEEP_TIME_IN_SECS) logging.warning( 'The checkpoint at %s could not be loaded after ' '%s seconds.', str(self._latest_export_dir), str(self._timeout)) return False