building MPNN (7a0ea258) · Commits · 钟慕尧 / deepchem

deepchem/data/datasets.py

+2 −1

Original line number	Diff line number	Diff line
		@@ -96,6 +96,7 @@ def pad_batch(batch_size, X_b, y_b, w_b, ids_b):

		# Fill in batch arrays
		start = 0
		w_out[start:start + num_samples] = w_b[:]
		while start < batch_size:
		num_left = batch_size - start
		if num_left < num_samples:
		@@ -104,9 +105,9 @@ def pad_batch(batch_size, X_b, y_b, w_b, ids_b):
		increment = num_samples
		X_out[start:start + increment] = X_b[:increment]
		y_out[start:start + increment] = y_b[:increment]
		w_out[start:start + increment] = w_b[:increment]
		ids_out[start:start + increment] = ids_b[:increment]
		start += increment

		return (X_out, y_out, w_out, ids_out)

deepchem/feat/graph_features.py

+11 −5

Original line number	Diff line number	Diff line
		@@ -287,18 +287,24 @@ class WeaveFeaturizer(Featurizer):

		name = ['weave_mol']

		def __init__(self, graph_distance=True):
		def __init__(self, graph_distance=True, explicit_H=None):
		# Set dtype
		self.graph_distance = graph_distance
		self.dtype = object
		self.check_H = False
		if explicit_H is None:
		self.explicit_H = False
		self.check_H = True

		def _featurize(self, mol):
		"""Encodes mol as a WeaveMol object."""
		# Atom features
		if self.graph_distance:
		idx_nodes = [(a.GetIdx(), atom_features(a)) for a in mol.GetAtoms()]
		else:
		idx_nodes = [(a.GetIdx(), atom_features(a, explicit_H=False)) for a in mol.GetAtoms()]
		if self.check_H and not self.explicit_H:
		for a in mol.GetAtoms():
		if a.GetSymbol() == 'H':
		self.explicit_H = True
		break
		idx_nodes = [(a.GetIdx(), atom_features(a, explicit_H=self.explicit_H)) for a in mol.GetAtoms()]
		idx_nodes.sort() # Sort by ind to ensure same order as rd_kit
		idx, nodes = list(zip(*idx_nodes))

deepchem/models/tensorgraph/graph_layers.py

+53 −0

Original line number	Diff line number	Diff line
		@@ -911,3 +911,56 @@ class GatedRecurrentUnit(object):
		z * inputs
		return h

		class SetGather(Layer):
		""" General class for MPNN """

		def __init__(self,
		M,
		batch_size,
		n_hidden=100,
		**kwargs):
		"""
		Parameters
		----------
		T: int
		Number of message passing steps
		message_fn: str, optional
		message function in the model
		update_fn: str, optional
		update function in the model
		n_hidden: int, optional
		number of hidden units in the passing phase
		"""

		self.M = M
		self.batch_size = batch_size
		self.n_hidden = n_hidden
		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)

		def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
		""" Perform T steps of message passing """
		if in_layers is None:
		in_layers = self.in_layers
		in_layers = convert_to_layers(in_layers)

		# Extract atom_features
		atom_features = in_layers[0].out_tensor
		atom_split = in_layers[1].out_tensor

		for i in range(self.M):
		q_expanded = tf.gather(self.state, 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]



		if set_tensors:
		self.variables = self.trainable_weights
		self.out_tensor = out_tensor
		return out_tensor
		No newline at end of file

deepchem/models/tensorgraph/models/graph_models.py

+8 −5

Original line number	Diff line number	Diff line
		@@ -5,8 +5,10 @@ import tensorflow as tf
		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
		from deepchem.models.tensorgraph.layers import Dense, Concat, SoftMax, SoftMaxCrossEntropy, GraphConv, BatchNorm, \
		Combine_AP, Separate_AP, DTNNEmbedding, DTNNStep, DTNNGather, DAGLayer, \
		DAGGather, DTNNExtract, MessagePassing
		from deepchem.models.tensorgraph.layers import Dense, Concat, SoftMax, \
		SoftMaxCrossEntropy, GraphConv, BatchNorm, \
		GraphPool, GraphGather, WeightedError, BatchNormalization, Stack
		from deepchem.models.tensorgraph.layers import L2Loss, Label, Weights, Feature
		from deepchem.models.tensorgraph.tensor_graph import TensorGraph
		@@ -678,10 +680,11 @@ class MPNNTensorGraph(TensorGraph):
		message_passing = MessagePassing(self.T,
		message_fn='enn',
		update_fn='gru',
		self.n_hidden,
		n_hidden=self.n_hidden,
		in_layers=[self.atom_features,
		self.pair_features,
		self.atom_to_pair])
		atom_embeddings = Dense(self.n_hidden, in_layers=[message_passing])



		@@ -690,7 +693,7 @@ class MPNNTensorGraph(TensorGraph):
		for task in range(self.n_tasks):
		if self.mode == "classification":
		classification = Dense(
		out_channels=2, activation_fn=None, in_layers=[weave_gather])
		out_channels=2, activation_fn=None, in_layers=[])
		softmax = SoftMax(in_layers=[classification])
		self.add_output(softmax)

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

		label = Label(shape=(None, 1))

Admin message