Merge pull request #1383 from peastman/atomicconv (cee8c837) · Commits · 钟慕尧 / deepchem

deepchem/models/tensorgraph/layers.py

+28 −28

Original line number	Diff line number	Diff line
		@@ -3839,32 +3839,32 @@ class AtomicConvolution(Layer):
		M = Nbrs.get_shape()[-1].value
		B = X.get_shape()[0].value

		D = self.distance_tensor(X, Nbrs, self.boxsize, B, N, M, d)
		R = self.distance_matrix(D)
		sym = []
		rsf_zeros = tf.zeros((B, N, M))
		for i, param in enumerate(self.radial_params):

		# Create the variables.
		if tfe.in_eager_mode():
		if not self._built:
		self.variables += self._create_radial_variables(*param)
		param_variables = self.variables[3 * i:3 * i + 3]
		self.variables += self._create_radial_variables()
		rc, rs, re = self.variables
		else:
		param_variables = self._create_radial_variables(*param)
		rc, rs, re = self._create_radial_variables()

		# We apply the radial pooling filter before atom type conv
		# to reduce computation
		rsf = self.radial_symmetry_function(R, *param_variables)
		# Compute the distances and radial symmetry functions.
		D = self.distance_tensor(X, Nbrs, self.boxsize, B, N, M, d)
		R = self.distance_matrix(D)
		R = tf.reshape(R, [1] + R.shape.as_list())
		rsf = self.radial_symmetry_function(R, rc, rs, re)

		if not self.atom_types:
		cond = tf.not_equal(Nbrs_Z, 0.0)
		sym.append(tf.reduce_sum(tf.where(cond, rsf, rsf_zeros), 2))
		cond = tf.to_float(tf.not_equal(Nbrs_Z, 0.0))
		cond = tf.reshape(cond, R.shape)
		layer = tf.reduce_sum(cond * rsf, 3)
		else:
		sym = []
		for j in range(len(self.atom_types)):
		cond = tf.equal(Nbrs_Z, self.atom_types[j])
		sym.append(tf.reduce_sum(tf.where(cond, rsf, rsf_zeros), 2))
		cond = tf.to_float(tf.equal(Nbrs_Z, self.atom_types[j]))
		cond = tf.reshape(cond, R.shape)
		sym.append(tf.reduce_sum(cond * rsf, 3))
		layer = tf.concat(sym, 0)

		layer = tf.stack(sym)
		layer = tf.transpose(layer, [1, 2, 0]) # (l, B, N) -> (B, N, l)
		m, v = tf.nn.moments(layer, axes=[0])
		out_tensor = tf.nn.batch_normalization(layer, m, v, None, None, 1e-3)
		@@ -3875,12 +3875,12 @@ class AtomicConvolution(Layer):
		self._built = True
		return out_tensor

		def _create_radial_variables(self, rc, rs, e):
		with tf.name_scope(None, "NbrRadialSymmetryFunction", [rc, rs, e]):
		rc = create_variable(rc)
		rs = create_variable(rs)
		e = create_variable(e)
		return (rc, rs, e)
		def _create_radial_variables(self):
		vars = []
		for i in range(3):
		val = np.array([p[i] for p in self.radial_params]).reshape((-1, 1, 1, 1))
		vars.append(create_variable(val, dtype=tf.float32))
		return vars

		def radial_symmetry_function(self, R, rc, rs, e):
		"""Calculates radial symmetry function.

deepchem/models/tensorgraph/tests/test_layers_eager.py

+15 −15

Original line number	Diff line number	Diff line
		@@ -22,8 +22,8 @@ class TestLayersEager(test_util.TensorFlowTestCase):
		filters = 3
		kernel_size = 2
		batch_size = 10
		input = np.random.rand(batch_size, width, in_channels).astype(
		np.float32)
		input = np.random.rand(batch_size, width,
		in_channels).astype(np.float32)
		layer = layers.Conv1D(filters, kernel_size)
		result = layer(input)
		self.assertEqual(result.shape[0], batch_size)
		@@ -170,8 +170,8 @@ class TestLayersEager(test_util.TensorFlowTestCase):
		n_hidden = 7
		in_channels = 4
		n_steps = 6
		input = np.random.rand(batch_size, n_steps, in_channels).astype(
		np.float32)
		input = np.random.rand(batch_size, n_steps,
		in_channels).astype(np.float32)
		layer = layers.GRU(n_hidden, batch_size)
		result, state = layer(input)
		assert result.shape == (batch_size, n_steps, n_hidden)
		@@ -203,8 +203,8 @@ class TestLayersEager(test_util.TensorFlowTestCase):
		n_hidden = 7
		in_channels = 4
		n_steps = 6
		input = np.random.rand(batch_size, n_steps, in_channels).astype(
		np.float32)
		input = np.random.rand(batch_size, n_steps,
		in_channels).astype(np.float32)
		layer = layers.LSTM(n_hidden, batch_size)
		result, state = layer(input)
		assert result.shape == (batch_size, n_steps, n_hidden)
		@@ -426,8 +426,8 @@ class TestLayersEager(test_util.TensorFlowTestCase):
		batch_size = 10
		n_features = 5
		logits = np.random.rand(batch_size, n_features).astype(np.float32)
		labels = np.random.randint(0, 2, (batch_size, n_features)).astype(
		np.float32)
		labels = np.random.randint(0, 2,
		(batch_size, n_features)).astype(np.float32)
		result = layers.SigmoidCrossEntropy()(labels, logits)
		expected = tf.nn.sigmoid_cross_entropy_with_logits(
		labels=labels, logits=logits)
		@@ -480,8 +480,8 @@ class TestLayersEager(test_util.TensorFlowTestCase):
		filters = 3
		kernel_size = 2
		batch_size = 10
		input = np.random.rand(batch_size, length, width, in_channels).astype(
		np.float32)
		input = np.random.rand(batch_size, length, width,
		in_channels).astype(np.float32)
		layer = layers.Conv2D(filters, kernel_size=kernel_size)
		result = layer(input)
		assert result.shape == (batch_size, length, width, filters)
		@@ -540,8 +540,8 @@ class TestLayersEager(test_util.TensorFlowTestCase):
		kernel_size = 2
		stride = 2
		batch_size = 10
		input = np.random.rand(batch_size, length, width, in_channels).astype(
		np.float32)
		input = np.random.rand(batch_size, length, width,
		in_channels).astype(np.float32)
		layer = layers.Conv2DTranspose(
		filters, kernel_size=kernel_size, stride=stride)
		result = layer(input)
		@@ -837,8 +837,8 @@ class TestLayersEager(test_util.TensorFlowTestCase):
		max_neighbors = 2
		dimensions = 3
		params = [[5.0, 2.0, 0.5], [10.0, 2.0, 0.5]]
		input1 = np.random.rand(batch_size, max_atoms, dimensions).astype(
		np.float32)
		input1 = np.random.rand(batch_size, max_atoms,
		dimensions).astype(np.float32)
		input2 = np.random.randint(
		max_atoms, size=(batch_size, max_atoms, max_neighbors))
		input3 = np.random.randint(
		@@ -846,7 +846,7 @@ class TestLayersEager(test_util.TensorFlowTestCase):
		layer = layers.AtomicConvolution(radial_params=params)
		result = layer(input1, input2, input3)
		assert result.shape == (batch_size, max_atoms, len(params))
		assert len(layer.variables) == 3 * len(params)
		assert len(layer.variables) == 3

		def test_alpha_share_layer(self):
		"""Test invoking AlphaShareLayer in eager mode."""

Admin message