Merge pull request #917 from lilleswing/config_proto (0404791b) · Commits · 钟慕尧 / deepchem

deepchem/models/tensorgraph/tensor_graph.py

+4 −2

Original line number	Diff line number	Diff line
		@@ -30,6 +30,7 @@ class TensorGraph(Model):
		use_queue=True,
		graph=None,
		learning_rate=0.001,
		configproto=None,
		**kwargs):
		"""
		Parameters
		@@ -51,7 +52,7 @@ class TensorGraph(Model):
		is created.
		learning_rate: float or LearningRateSchedule
		the learning rate to use for optimization
		kwargs
		configproto: a tf.ConfigProto() object used to create tf.Session()
		"""

		# Layer Management
		@@ -66,6 +67,7 @@ class TensorGraph(Model):
		self.queue_installed = False
		self.optimizer = Adam(
		learning_rate=learning_rate, beta1=0.9, beta2=0.999, epsilon=1e-7)
		self.configproto = configproto

		# Singular place to hold Tensor objects which don't serialize
		# These have to be reconstructed on restoring from pickle
		@@ -470,7 +472,7 @@ class TensorGraph(Model):
		self.rnn_final_states += layer.rnn_final_states
		self.rnn_zero_states += layer.rnn_zero_states
		layer.add_summary_to_tg()
		self.session = tf.Session()
		self.session = tf.Session(config=self.configproto)
		self.built = True

		# Ensure all training operators have been created.