def test_save_graph_model_default_session(self): x = tf.placeholder(shape=[None, 10], dtype=tf.float32, name='inp') weights = tf.constant(1., shape=(10, 2), name='weights') model_path = os.path.join(tf.test.get_temp_dir(), 'default') utils.save_graph_model( tf.Session(), model_path, {'x': x}, {'w': weights}, ['tag']) self.assertTrue(os.path.isfile(os.path.join(model_path, 'saved_model.pb')))
def save_model_with_metadata(self, file_path): """Saves the model and the generated metadata to the given file path. Args: file_path: Path to save the model and the metadata. It can be a GCS bucket or a local folder. The folder needs to be empty. Returns: Full file path where the model and the metadata are written. """ md_dict = self.get_metadata() if not self._serving_inputs: self._serving_inputs = self._build_input_signature( self._inputs, self._session.graph) if not self._serving_outputs: self._serving_outputs = self._build_output_signature( self._outputs, self._session.graph) utils.save_graph_model(self._session, file_path, self._serving_inputs, self._serving_outputs, self._tags, **self._saved_model_args) common_utils.write_metadata_to_file(md_dict, file_path) return file_path
def test_save_graph_model_kwargs(self): x = tf.placeholder(shape=[None, 10], dtype=tf.float32, name='inp') weights = tf.constant(1., shape=(10, 2), name='weights') model_path = os.path.join(tf.test.get_temp_dir(), 'kwargs') utils.save_graph_model( tf.Session(), model_path, {'x': x}, {'w': weights}, ['tag'], main_op=tf.tables_initializer(), strip_default_attrs=False) self.assertTrue(os.path.isfile(os.path.join(model_path, 'saved_model.pb')))
def test_save_graph_model_explicit_session(self): sess = tf.Session(graph=tf.Graph()) with sess.graph.as_default(): x = tf.placeholder(shape=[None, 10], dtype=tf.float32, name='inp') weights = tf.constant(1., shape=(10, 2), name='weights') model_path = os.path.join(tf.test.get_temp_dir(), 'explicit') utils.save_graph_model(sess, model_path, {'x': x}, {'w': weights}, ['tag']) self.assertTrue(os.path.isfile(os.path.join(model_path, 'saved_model.pb'))) tf.reset_default_graph() loading_session = tf.Session(graph=tf.Graph()) with loading_session.graph.as_default(): tf.saved_model.loader.load(loading_session, ['tag'], model_path) self.assertIn(x.op.name, [n.name for n in loading_session.graph.as_graph_def().node]) self.assertIn(weights.op.name, [n.name for n in loading_session.graph.as_graph_def().node])