first build of MPNN (8f2e61dc) · Commits · 钟慕尧 / deepchem

deepchem/models/init.py

+1 −1

Original line number	Diff line number	Diff line
		@@ -28,5 +28,5 @@ from deepchem.models.tensorflow_models.progressive_multitask import ProgressiveM
		from deepchem.models.tensorflow_models.progressive_joint import ProgressiveJointRegressor
		from deepchem.models.tensorflow_models.IRV import TensorflowMultiTaskIRVClassifier
		from deepchem.models.tensorgraph.tensor_graph import TensorGraph
		from deepchem.models.tensorgraph.models.graph_models import WeaveTensorGraph, DTNNTensorGraph, DAGTensorGraph, GraphConvTensorGraph
		from deepchem.models.tensorgraph.models.graph_models import WeaveTensorGraph, DTNNTensorGraph, DAGTensorGraph, GraphConvTensorGraph, MPNNTensorGraph
		from deepchem.models.tensorgraph.models.symmetry_function_regression import BPSymmetryFunctionRegression, ANIRegression

deepchem/models/tensorgraph/graph_layers.py

+31 −7

Original line number	Diff line number	Diff line
		@@ -918,6 +918,7 @@ class SetGather(Layer):
		M,
		batch_size,
		n_hidden=100,
		init='orthogonal',
		**kwargs):
		"""
		Parameters
		@@ -935,12 +936,16 @@ class SetGather(Layer):
		self.M = M
		self.batch_size = batch_size
		self.n_hidden = n_hidden
		self.init = initializations.get(init)
		super(SetGather, self).__init__(**kwargs)

		def build(self, pair_features, n_pair_features):
		self.state = tf.zeros((self.batch_size, self.n_hidden))
		self.lstm_cell = tf.contrib.rnn.BasicLSTMCell(self.n_hidden,
		state_is_tuple=False)
		self.U = self.init((2self.n_hidden, 4self.n_hidden))
		self.b = tf.Variable(
		np.concatenate((np.zeros(self.n_hidden), np.ones(self.n_hidden),
		np.zeros(self.n_hidden), np.zeros(self.n_hidden))),
		dtype=tf.float32)


		def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
		""" Perform T steps of message passing """
		@@ -948,19 +953,38 @@ class SetGather(Layer):
		in_layers = self.in_layers
		in_layers = convert_to_layers(in_layers)

		self.build()
		# Extract atom_features
		atom_features = in_layers[0].out_tensor
		atom_split = in_layers[1].out_tensor

		c = tf.zeros((self.batch_size, self.n_hidden))
		h = tf.zeros((self.batch_size, self.n_hidden))

		for i in range(self.M):
		q_expanded = tf.gather(self.state, atom_split)
		q_expanded = tf.gather(h, atom_split)
		e = tf.reduce_sum(atom_features * q_expanded, 1)
		e_mols = tf.dynamic_partition(e, atom_split, self.batch_size)
		a = [tf.nn.softmax(e_mol) for e_mol in e_mols]


		a = tf.concat([tf.nn.softmax(e_mol) for e_mol in e_mols], 0)
		r = tf.segment_sum(tf.reshape(a, [-1,1]) * atom_features, atom_split)
		q_star = tf.concat([h, r], axis=1)
		h, c = self.LSTMStep(q_star, c)

		out_tensor = q_star
		if set_tensors:
		self.variables = self.trainable_weights
		self.out_tensor = out_tensor
		return out_tensor

		def LSTMStep(self, h, c, x=None):

		# Taken from Keras code [citation needed]
		z = tf.nn.xw_plus_b(h, self.U, self.b)
		i = tf.nn.sigmoid(z[:, :self.n_hidden])
		f = tf.nn.sigmoid(z[:, self.n_hidden:2 * self.n_hidden])
		o = tf.nn.sigmoid(z[:, 2 * self.n_hidden:3 * self.n_hidden])
		z3 = z[:, 3 * self.n_hidden:]
		c_out = f * c + i * tf.nn.tanh(z3)
		h_out = o * tf.nn.tanh(c_out)

		return h_out, c_out
		No newline at end of file

deepchem/models/tensorgraph/models/graph_models.py

+14 −10

Original line number	Diff line number	Diff line
		@@ -6,7 +6,7 @@ from deepchem.feat.mol_graphs import ConvMol
		from deepchem.metrics import to_one_hot, from_one_hot
		from deepchem.models.tensorgraph.graph_layers import WeaveLayer, WeaveGather, \
		Combine_AP, Separate_AP, DTNNEmbedding, DTNNStep, DTNNGather, DAGLayer, \
		DAGGather, DTNNExtract, MessagePassing
		DAGGather, DTNNExtract, MessagePassing, SetGather
		from deepchem.models.tensorgraph.layers import Dense, Concat, SoftMax, \
		SoftMaxCrossEntropy, GraphConv, BatchNorm, \
		GraphPool, GraphGather, WeightedError, Dropout, BatchNormalization, Stack
		@@ -684,10 +684,12 @@ class MPNNTensorGraph(TensorGraph):

		def __init__(self,
		n_tasks,
		batch_size,
		n_atom_feat=70,
		n_pair_feat=8,
		n_hidden=100,
		T=5,
		M=10,
		**kwargs):
		"""
		Parameters
		@@ -705,20 +707,18 @@ class MPNNTensorGraph(TensorGraph):

		"""
		self.n_tasks = n_tasks
		self.batch_size = batch_size
		self.n_atom_feat = n_atom_feat
		self.n_pair_feat = n_pair_feat
		self.n_hidden = n_hidden
		self.T = T
		self.M = M
		super(MPNNTensorGraph, self).__init__(**kwargs)
		self.build_graph()

		def build_graph(self):
		"""Building graph structures:
		Features => WeaveLayer => WeaveLayer => Dense => WeaveGather => Classification or Regression
		"""
		self.atom_features = Feature(shape=(None, self.n_atom_feat))
		self.pair_features = Feature(shape=(None, self.n_pair_feat))
		self.pair_split = Feature(shape=(None,), dtype=tf.int32)
		self.atom_split = Feature(shape=(None,), dtype=tf.int32)
		self.atom_to_pair = Feature(shape=(None, 2), dtype=tf.int32)

		@@ -730,15 +730,20 @@ class MPNNTensorGraph(TensorGraph):
		self.pair_features,
		self.atom_to_pair])
		atom_embeddings = Dense(self.n_hidden, in_layers=[message_passing])
		mol_embeddings = SetGather(self.M,
		self.batch_size,
		n_hidden=self.n_hidden,
		in_layers=[atom_embeddings, self.atom_split])



		dense1 = Dense(out_channels=2*self.n_hidden,
		activation_fn=tf.nn.relu,
		in_layers=[mol_embeddings])
		costs = []
		self.labels_fd = []
		for task in range(self.n_tasks):
		if self.mode == "classification":
		classification = Dense(
		out_channels=2, activation_fn=None, in_layers=[])
		out_channels=2, activation_fn=None, in_layers=[dense1])
		softmax = SoftMax(in_layers=[classification])
		self.add_output(softmax)

		@@ -748,7 +753,7 @@ class MPNNTensorGraph(TensorGraph):
		costs.append(cost)
		if self.mode == "regression":
		regression = Dense(
		out_channels=1, activation_fn=None, in_layers=[])
		out_channels=1, activation_fn=None, in_layers=[dense1])
		self.add_output(regression)

		label = Label(shape=(None, 1))
		@@ -818,7 +823,6 @@ class MPNNTensorGraph(TensorGraph):

		feed_dict[self.atom_features] = np.concatenate(atom_feat, axis=0)
		feed_dict[self.pair_features] = np.concatenate(pair_feat, axis=0)
		feed_dict[self.pair_split] = np.array(pair_split)
		feed_dict[self.atom_split] = np.array(atom_split)
		feed_dict[self.atom_to_pair] = np.concatenate(atom_to_pair, axis=0)
		yield feed_dict
		No newline at end of file

Admin message