Continuing support for eager mode in TensorGraph (db395e85) · Commits · 钟慕尧 / deepchem

deepchem/models/tensorgraph/layers.py

+18 −13

Original line number	Diff line number	Diff line
		@@ -38,7 +38,7 @@ class Layer(object):
		if tfe.in_eager_mode():
		self.variables = []
		self._built = False
		self._non_pickle_fields = ['variables']
		self._non_pickle_fields = ['variables', '_built']
		else:
		self.variable_scope = ''
		self._non_pickle_fields = [
		@@ -98,6 +98,8 @@ class Layer(object):
		-------
		Layer
		"""
		if tfe.in_eager_mode():
		raise ValueError('shared() is not supported in eager mode')
		if self.variable_scope == '':
		return self.clone(in_layers)
		raise ValueError('%s does not implement shared()' % self.__class__.__name__)
		@@ -262,8 +264,6 @@ class Layer(object):
		This means the newly created layers will share variables with the original
		ones.
		"""
		if tfe.in_eager_mode():
		raise ValueError('copy() is not supported in eager mode')
		if self in replacements:
		return replacements[self]
		copied_inputs = [
		@@ -280,6 +280,9 @@ class Layer(object):
		variables = variables_graph.get_layer_variables(self)
		if len(variables) > 0:
		with variables_graph._get_tf("Graph").as_default():
		if tfe.in_eager_mode():
		values = [v.numpy() for v in variables]
		else:
		values = variables_graph.session.run(variables)
		copy.set_variable_initial_values(values)
		return copy
		@@ -381,6 +384,8 @@ class SharedVariableScope(Layer):
		self._shared_with = None

		def shared(self, in_layers):
		if tfe.in_eager_mode():
		raise ValueError('shared() is not supported in eager mode')
		copy = self.clone(in_layers)
		self._reuse = True
		copy._reuse = True
		@@ -1753,11 +1758,11 @@ class ReduceMean(Layer):
		def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
		inputs = self._get_input_tensors(in_layers)
		if len(inputs) > 1:
		self.out_tensor = tf.stack(inputs)
		out_tensor = tf.stack(inputs)
		else:
		self.out_tensor = inputs[0]
		out_tensor = inputs[0]

		out_tensor = tf.reduce_mean(self.out_tensor, axis=self.axis)
		out_tensor = tf.reduce_mean(out_tensor, axis=self.axis)
		if set_tensors:
		self.out_tensor = out_tensor
		return out_tensor
		@@ -1784,11 +1789,11 @@ class ReduceMax(Layer):
		def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
		inputs = self._get_input_tensors(in_layers)
		if len(inputs) > 1:
		self.out_tensor = tf.stack(inputs)
		out_tensor = tf.stack(inputs)
		else:
		self.out_tensor = inputs[0]
		out_tensor = inputs[0]

		out_tensor = tf.reduce_max(self.out_tensor, axis=self.axis)
		out_tensor = tf.reduce_max(out_tensor, axis=self.axis)
		if set_tensors:
		self.out_tensor = out_tensor
		return out_tensor
		@@ -1834,11 +1839,11 @@ class ReduceSum(Layer):
		def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
		inputs = self._get_input_tensors(in_layers)
		if len(inputs) > 1:
		self.out_tensor = tf.stack(inputs)
		out_tensor = tf.stack(inputs)
		else:
		self.out_tensor = inputs[0]
		out_tensor = inputs[0]

		out_tensor = tf.reduce_sum(self.out_tensor, axis=self.axis)
		out_tensor = tf.reduce_sum(out_tensor, axis=self.axis)
		if set_tensors:
		self.out_tensor = out_tensor
		return out_tensor

deepchem/models/tensorgraph/tensor_graph.py

+33 −7

Original line number	Diff line number	Diff line
		@@ -191,10 +191,16 @@ class TensorGraph(Model):
		# In eager mode we want an optimizer and a function to compute the
		# gradient of the loss.

		submodel_vars = None
		if submodel is None:
		optimizer = self._get_tf("Optimizer")
		else:
		optimizer = submodel.create_optimizer()
		if submodel.layers is not None:
		submodel_vars = set()
		for layer in submodel.layers:
		for var in layer.variables:
		submodel_vars.add(var)
		val_grad_fn = tfe.implicit_value_and_gradients(
		lambda x: self._run_graph([loss], x, True)[0])
		else:
		@@ -237,6 +243,10 @@ class TensorGraph(Model):
		n_samples % self.tensorboard_log_frequency == 0)
		if tfe.in_eager_mode():
		value, grads_and_vars = val_grad_fn(feed_dict)
		if submodel_vars is not None:
		grads_and_vars = [
		x for x in grads_and_vars if x[1] in submodel_vars
		]
		optimizer.apply_gradients(grads_and_vars)
		avg_loss += value
		else:
		@@ -362,6 +372,8 @@ class TensorGraph(Model):
		break
		n_samples += 1
		feed_results = self._run_graph(outputs, feed_dict, False)
		if tfe.in_eager_mode():
		feed_results = [f.numpy() for f in feed_results]
		if len(feed_results) > 1:
		if len(transformers):
		raise ValueError("Does not support transformations "
		@@ -503,17 +515,27 @@ class TensorGraph(Model):
		shape = [1 if s is None else s for s in layer.shape]
		tensor = tf.zeros(shape, layer.dtype)
		else:
		with tf.name_scope(layer.name):
		tensor = layer.create_tensor(in_layers=inputs, set_tensors=False)
		tensors[layer] = tensor
		return tensor

		tensors = {}
		with self._get_tf("Graph").as_default():
		# Build the layers.

		build_layers(self.loss, tensors)
		for output in self.outputs:
		build_layers(output, tensors)
		for submodel in self.submodels:
		build_layers(submodel.loss, tensors)

		# Initialize variables.

		for layer in self.layers.values():
		if layer.variable_values is not None:
		for var, val in zip(layer.variables, layer.variable_values):
		var.assign(val)
		self.session = None
		self._training_placeholder = None
		self.built = True
		@@ -861,12 +883,14 @@ class TensorGraph(Model):
		with self._get_tf("Graph").as_default():
		reader = NewCheckpointReader(checkpoint)
		var_names = set([x for x in reader.get_variable_to_shape_map()])
		var_map = {
		x.op.name: x
		for x in tf.global_variables()
		if x.op.name in var_names
		}
		saver = tf.train.Saver(var_list=var_map)
		var_list = []
		for var in self.get_variables():
		name = var.name
		if ':' in name:
		name = name[:name.rfind(':')]
		if name in var_names:
		var_list.append(var)
		saver = tf.train.Saver(var_list=var_list)
		saver.restore(self.session, checkpoint)

		def get_num_tasks(self):
		@@ -992,6 +1016,8 @@ class TensorGraph(Model):
		"""
		# Check the inputs.

		if tfe.in_eager_mode():
		raise ValueError('make_estimator() is not supported in eager mode')
		if len(feature_columns) != len(self.features):
		raise ValueError(
		'This model requires %d feature column(s)' % len(self.features))

deepchem/models/tensorgraph/tests/test_tensor_graph.py

+65 −1

Original line number	Diff line number	Diff line
		@@ -16,6 +16,8 @@ from deepchem.models.tensorgraph.layers import Feature, Label
		from deepchem.models.tensorgraph.layers import ReduceSquareDifference, Add
		from deepchem.models.tensorgraph.tensor_graph import TensorGraph
		from deepchem.models.tensorgraph.optimizers import GradientDescent, ExponentialDecay
		import tensorflow.contrib.eager as tfe
		from tensorflow.python.eager import context


		class TestTensorGraph(unittest.TestCase):
		@@ -42,6 +44,11 @@ class TestTensorGraph(unittest.TestCase):
		prediction = np.squeeze(tg.predict_on_batch(X))
		assert_true(np.all(np.isclose(prediction, y, atol=0.4)))

		def test_single_task_classifier_eager(self):
		with context.eager_mode():
		with tfe.IsolateTest():
		self.test_single_task_classifier()

		@flaky
		def test_multi_task_classifier(self):
		n_data_points = 20
		@@ -86,6 +93,12 @@ class TestTensorGraph(unittest.TestCase):
		y_pred = predictions[i]
		assert_true(np.all(np.isclose(y_pred, y_real, atol=0.6)))

		@flaky
		def test_multi_task_classifier_eager(self):
		with context.eager_mode():
		with tfe.IsolateTest():
		self.test_multi_task_classifier()

		def test_single_task_regressor(self):
		n_data_points = 20
		n_features = 2
		@@ -103,6 +116,11 @@ class TestTensorGraph(unittest.TestCase):
		prediction = np.squeeze(tg.predict_on_batch(X))
		assert_true(np.all(np.isclose(prediction, y, atol=3.0)))

		def test_single_task_regressor_eager(self):
		with context.eager_mode():
		with tfe.IsolateTest():
		self.test_single_task_regressor()

		def test_multi_task_regressor(self):
		n_data_points = 20
		n_features = 2
		@@ -145,6 +163,11 @@ class TestTensorGraph(unittest.TestCase):
		y_pred = predictions[i]
		assert_true(np.all(np.isclose(y_pred, y_real, atol=1.5)))

		def test_multi_task_regressor_eager(self):
		with context.eager_mode():
		with tfe.IsolateTest():
		self.test_multi_task_regressor()

		@flaky
		def test_no_queue(self):
		n_data_points = 20
		@@ -193,6 +216,12 @@ class TestTensorGraph(unittest.TestCase):
		prediction2 = np.squeeze(tg1.predict_on_batch(X))
		assert_true(np.all(np.isclose(prediction, prediction2, atol=0.01)))

		@flaky
		def test_set_optimizer_eager(self):
		with context.eager_mode():
		with tfe.IsolateTest():
		self.test_set_optimizer()

		def test_tensorboard(self):
		n_data_points = 20
		n_features = 2
		@@ -221,6 +250,11 @@ class TestTensorGraph(unittest.TestCase):
		file_size = os.stat(event_file).st_size
		assert_true(file_size > 0)

		def test_tensorboard_eager(self):
		with context.eager_mode():
		with tfe.IsolateTest():
		self.test_tensorboard()

		def test_save_load(self):
		n_data_points = 20
		n_features = 2
		@@ -248,6 +282,11 @@ class TestTensorGraph(unittest.TestCase):
		prediction2 = np.squeeze(tg1.predict_on_batch(X))
		assert_true(np.all(np.isclose(prediction, prediction2, atol=0.01)))

		def test_save_load_eager(self):
		with context.eager_mode():
		with tfe.IsolateTest():
		self.test_save_load()

		def test_shared_layer(self):
		n_data_points = 20
		n_features = 2
		@@ -325,6 +364,11 @@ class TestTensorGraph(unittest.TestCase):
		value = tg.predict_on_batch(np.array([0]), outputs=o)
		assert np.array_equal(e, value)

		def test_operators_eager(self):
		with context.eager_mode():
		with tfe.IsolateTest():
		self.test_operators()

		def test_initialize_variable(self):
		"""Test methods for initializing a variable."""
		# Set by variable constructor.
		@@ -339,11 +383,18 @@ class TestTensorGraph(unittest.TestCase):
		# Set by set_variable_initial_values().

		tg = dc.models.TensorGraph(use_queue=False)
		features = Feature(shape=(None, 1))
		tg.set_loss(Dense(1, in_layers=features))
		var = Variable([10.0])
		var.set_variable_initial_values([[15.0]])
		tg.add_output(var)
		assert tg.predict_on_batch(np.zeros((1, 1))) == [15.0]

		def test_initialize_variable_eager(self):
		with context.eager_mode():
		with tfe.IsolateTest():
		self.test_initialize_variable()

		def test_copy_layers(self):
		"""Test copying layers."""
		tg = dc.models.TensorGraph()
		@@ -363,10 +414,18 @@ class TestTensorGraph(unittest.TestCase):
		assert copy.in_layers[1] == replacements[constant]
		variables = tg.get_layer_variables(dense)
		with tg._get_tf("Graph").as_default():
		if tfe.in_eager_mode():
		values = [v.numpy() for v in variables]
		else:
		values = tg.session.run(variables)
		for v1, v2 in zip(values, copy.in_layers[0].variable_values):
		assert np.array_equal(v1, v2)

		def test_copy_layers_eager(self):
		with context.eager_mode():
		with tfe.IsolateTest():
		self.test_copy_layers()

		def test_copy_layers_shared(self):
		"""Test copying layers with shared variables."""
		tg = dc.models.TensorGraph()
		@@ -432,3 +491,8 @@ class TestTensorGraph(unittest.TestCase):
		1.0, tg.predict_on_batch(data, outputs=var1)[0], places=4)
		self.assertAlmostEqual(
		0.0, tg.predict_on_batch(data, outputs=var2)[0], places=4)

		def test_submodels_eager(self):
		with context.eager_mode():
		with tfe.IsolateTest():
		self.test_submodels()

Admin message