Commit bec166dd authored by Yuxin Wu's avatar Yuxin Wu

Use NewSessionCreate in PredictConfig, so that variables are initialized (fix #1225)

parent 7c8cbba0
...@@ -8,6 +8,7 @@ from ..compat import tfv1 as tf ...@@ -8,6 +8,7 @@ from ..compat import tfv1 as tf
from ..graph_builder import ModelDescBase from ..graph_builder import ModelDescBase
from ..tfutils import get_default_sess_config from ..tfutils import get_default_sess_config
from ..tfutils.sessinit import JustCurrentSession, SessionInit from ..tfutils.sessinit import JustCurrentSession, SessionInit
from ..tfutils.sesscreate import NewSessionCreator
from ..tfutils.tower import TowerFuncWrapper from ..tfutils.tower import TowerFuncWrapper
from ..utils import logger from ..utils import logger
...@@ -61,7 +62,7 @@ class PredictConfig(object): ...@@ -61,7 +62,7 @@ class PredictConfig(object):
tensors can be any tensor in the graph that's computable from the tensors correponding to `input_names`. tensors can be any tensor in the graph that's computable from the tensors correponding to `input_names`.
session_creator (tf.train.SessionCreator): how to create the session_creator (tf.train.SessionCreator): how to create the
session. Defaults to :class:`tf.train.ChiefSessionCreator()`. session. Defaults to :class:`NewSessionCreator()`.
session_init (SessionInit): how to initialize variables of the session. session_init (SessionInit): how to initialize variables of the session.
Defaults to do nothing. Defaults to do nothing.
...@@ -98,7 +99,7 @@ class PredictConfig(object): ...@@ -98,7 +99,7 @@ class PredictConfig(object):
assert_type(self.session_init, SessionInit, 'session_init') assert_type(self.session_init, SessionInit, 'session_init')
if session_creator is None: if session_creator is None:
self.session_creator = tf.train.ChiefSessionCreator(config=get_default_sess_config()) self.session_creator = NewSessionCreator(config=get_default_sess_config())
else: else:
self.session_creator = session_creator self.session_creator = session_creator
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment