Continuing to refactor graph models (a453eb76) · Commits · 钟慕尧 / deepchem

deepchem/models/tensorgraph/models/graph_models.py

+74 −48

Original line number	Diff line number	Diff line
		@@ -395,6 +395,8 @@ class DAGModel(TensorGraph):
		layer_sizes_gather=[100],
		dropout=None,
		mode="classification",
		n_classes=2,
		uncertainty=False,
		**kwargs):
		"""
		Parameters
		@@ -419,6 +421,11 @@ class DAGModel(TensorGraph):
		Dropout probability, applied after each propagation step and gather step.
		mode: str, optional
		Either "classification" or "regression" for type of model.
		n_classes: int
		the number of classes to predict (only used in classification mode)
		uncertainty: bool
		if True, include extra outputs and loss terms to enable the uncertainty
		in outputs to be predicted
		"""
		self.n_tasks = n_tasks
		self.max_atoms = max_atoms
		@@ -429,6 +436,13 @@ class DAGModel(TensorGraph):
		self.layer_sizes_gather = layer_sizes_gather
		self.dropout = dropout
		self.mode = mode
		self.n_classes = n_classes
		self.uncertainty = uncertainty
		if uncertainty:
		if mode != "regression":
		raise ValueError("Uncertainty is only supported in regression mode")
		if dropout == 0.0:
		raise ValueError('Dropout must be included to predict uncertainty')
		super(DAGModel, self).__init__(**kwargs)
		self.build_graph()

		@@ -464,35 +478,39 @@ class DAGModel(TensorGraph):
		dropout=self.dropout,
		in_layers=[dag_layer1, self.membership])

		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=[dag_gather])
		softmax = SoftMax(in_layers=[classification])
		self.add_output(softmax)

		label = Label(shape=(None, 2))
		self.labels_fd.append(label)
		cost = SoftMaxCrossEntropy(in_layers=[label, classification])
		costs.append(cost)
		if self.mode == "regression":
		regression = Dense(
		out_channels=1, activation_fn=None, in_layers=[dag_gather])
		self.add_output(regression)

		label = Label(shape=(None, 1))
		self.labels_fd.append(label)
		cost = L2Loss(in_layers=[label, regression])
		costs.append(cost)
		if self.mode == "classification":
		all_cost = Stack(in_layers=costs, axis=1)
		elif self.mode == "regression":
		all_cost = Stack(in_layers=costs, axis=1)
		self.weights = Weights(shape=(None, self.n_tasks))
		loss = WeightedError(in_layers=[all_cost, self.weights])
		self.set_loss(loss)
		n_tasks = self.n_tasks
		weights = Weights(shape=(None, n_tasks))
		if self.mode == 'classification':
		n_classes = self.n_classes
		labels = Label(shape=(None, n_tasks, n_classes))
		logits = Reshape(
		shape=(None, n_tasks, n_classes),
		in_layers=[
		Dense(in_layers=dag_gather, out_channels=n_tasks * n_classes)
		])
		output = SoftMax(logits)
		self.add_output(output)
		loss = SoftMaxCrossEntropy(in_layers=[labels, logits])
		weighted_loss = WeightedError(in_layers=[loss, weights])
		self.set_loss(weighted_loss)
		else:
		labels = Label(shape=(None, n_tasks))
		output = Reshape(
		shape=(None, n_tasks),
		in_layers=[Dense(in_layers=dag_gather, out_channels=n_tasks)])
		self.add_output(output)
		if self.uncertainty:
		log_var = Reshape(
		shape=(None, n_tasks),
		in_layers=[Dense(in_layers=dag_gather, out_channels=n_tasks)])
		var = Exp(log_var)
		self.add_variance(var)
		diff = labels - output
		weighted_loss = weights * (diff * diff / var + log_var)
		weighted_loss = ReduceSum(ReduceMean(weighted_loss, axis=[1]))
		else:
		weighted_loss = ReduceSum(L2Loss(in_layers=[labels, output, weights]))
		self.set_loss(weighted_loss)

		def default_generator(self,
		dataset,
		@@ -502,8 +520,6 @@ class DAGModel(TensorGraph):
		pad_batches=True):
		"""TensorGraph style implementation"""
		for epoch in range(epochs):
		if not predict:
		print('Starting epoch %i' % epoch)
		for (X_b, y_b, w_b, ids_b) in dataset.iterbatches(
		batch_size=self.batch_size,
		deterministic=deterministic,
		@@ -511,13 +527,15 @@ class DAGModel(TensorGraph):

		feed_dict = dict()
		if y_b is not None:
		for index, label in enumerate(self.labels_fd):
		if self.mode == "classification":
		feed_dict[label] = to_one_hot(y_b[:, index])
		if self.mode == "regression":
		feed_dict[label] = y_b[:, index:index + 1]
		if self.mode == 'classification':
		feed_dict[self.labels[0]] = to_one_hot(y_b.flatten(),
		self.n_classes).reshape(
		-1, self.n_tasks,
		self.n_classes)
		else:
		feed_dict[self.labels[0]] = y_b
		if w_b is not None:
		feed_dict[self.weights] = w_b
		feed_dict[self.task_weights[0]] = w_b

		atoms_per_mol = [mol.get_num_atoms() for mol in X_b]
		n_atoms = sum(atoms_per_mol)
		@@ -568,7 +586,7 @@ class GraphConvModel(TensorGraph):
		n_tasks,
		graph_conv_layers=[64, 64],
		dense_layer_size=128,
		dropout=0.0,
		dropouts=0.0,
		mode="classification",
		number_atom_features=75,
		n_classes=2,
		@@ -583,8 +601,11 @@ class GraphConvModel(TensorGraph):
		Width of channels for the Graph Convolution Layers
		dense_layer_size: int
		Width of channels for Atom Level Dense Layer before GraphPool
		dropout: float
		Droupout dropout probability. Dropout is applied after the per Atom Level Dense Layer
		dropouts: list or float
		the dropout probablity to use for each layer. The length of this list should equal
		len(graph_conv_layers)+1 (one value for each convolution layer, and one for the
		dense layer). Alternatively this may be a single value instead of a list, in which
		case the same value is used for every layer.
		mode: str
		Either "classification" or "regression"
		number_atom_features: int
		@@ -602,17 +623,22 @@ class GraphConvModel(TensorGraph):
		self.n_tasks = n_tasks
		self.mode = mode
		self.dense_layer_size = dense_layer_size
		self.dropout = dropout
		self.graph_conv_layers = graph_conv_layers
		kwargs['use_queue'] = False
		self.number_atom_features = number_atom_features
		self.n_classes = n_classes
		self.uncertainty = uncertainty
		if not isinstance(dropouts, collections.Sequence):
		dropouts = [dropouts] * (len(graph_conv_layers) + 1)
		if len(dropouts) != len(graph_conv_layers) + 1:
		raise ValueError('Wrong number of dropout probabilities provided')
		self.dropouts = dropouts
		if uncertainty:
		if mode != "regression":
		raise ValueError("Uncertainty is only supported in regression mode")
		if dropout == 0.0:
		raise ValueError('Dropout must be included to predict uncertainty')
		if any(d == 0.0 for d in dropouts):
		raise ValueError(
		'Dropout must be included in every layer to predict uncertainty')
		super(GraphConvModel, self).__init__(**kwargs)
		self.build_graph()

		@@ -629,12 +655,12 @@ class GraphConvModel(TensorGraph):
		deg_adj = Feature(shape=(None, i + 1), dtype=tf.int32)
		self.deg_adjs.append(deg_adj)
		in_layer = self.atom_features
		for layer_size in self.graph_conv_layers:
		for layer_size, dropout in zip(self.graph_conv_layers, self.dropouts):
		gc1_in = [in_layer, self.degree_slice, self.membership] + self.deg_adjs
		gc1 = GraphConv(layer_size, activation_fn=tf.nn.relu, in_layers=gc1_in)
		batch_norm1 = BatchNorm(in_layers=[gc1])
		if self.dropout > 0.0:
		batch_norm1 = Dropout(self.dropout, in_layers=batch_norm1)
		if dropout > 0.0:
		batch_norm1 = Dropout(dropout, in_layers=batch_norm1)
		gp_in = [batch_norm1, self.degree_slice, self.membership] + self.deg_adjs
		in_layer = GraphPool(in_layers=gp_in)
		dense = Dense(
		@@ -642,8 +668,8 @@ class GraphConvModel(TensorGraph):
		activation_fn=tf.nn.relu,
		in_layers=[in_layer])
		batch_norm3 = BatchNorm(in_layers=[dense])
		if self.dropout > 0.0:
		batch_norm3 = Dropout(self.dropout, in_layers=batch_norm3)
		if self.dropouts[-1] > 0.0:
		batch_norm3 = Dropout(self.dropouts[-1], in_layers=batch_norm3)
		readout = GraphGather(
		batch_size=self.batch_size,
		activation_fn=tf.nn.tanh,

deepchem/models/tensorgraph/models/test_graph_models.py

+85 −12

Original line number	Diff line number	Diff line
		@@ -2,14 +2,15 @@ import unittest

		import numpy as np

		import deepchem
		import deepchem as dc
		from deepchem.data import NumpyDataset
		from deepchem.models import GraphConvModel
		from deepchem.models import GraphConvModel, DAGModel
		from deepchem.models import TensorGraph
		from deepchem.molnet import load_bace_classification, load_delaney
		from deepchem.models.tensorgraph.layers import ReduceSum, L2Loss
		from deepchem.models import WeaveModel
		from deepchem.feat import ConvMolFeaturizer
		from nose.plugins.attrib import attr


		class TestGraphModels(unittest.TestCase):
		@@ -31,13 +32,12 @@ class TestGraphModels(unittest.TestCase):

		if mode == 'classification':
		y = np.random.randint(0, 2, size=(data_points, len(tasks)))
		metric = deepchem.metrics.Metric(
		deepchem.metrics.roc_auc_score, np.mean, mode="classification")
		transformers = []
		metric = dc.metrics.Metric(
		dc.metrics.roc_auc_score, np.mean, mode="classification")
		else:
		y = np.random.normal(size=(data_points, len(tasks)))
		metric = deepchem.metrics.Metric(
		deepchem.metrics.mean_absolute_error, mode="regression")
		metric = dc.metrics.Metric(
		dc.metrics.mean_absolute_error, mode="regression")

		ds = NumpyDataset(train.X[:10], y, w, train.ids[:10])

		@@ -87,19 +87,16 @@ class TestGraphModels(unittest.TestCase):
		len(tasks),
		batch_size=batch_size,
		mode='regression',
		dropout=0.1,
		dropouts=0.1,
		uncertainty=True)

		model.fit(dataset, nb_epoch=100)
		scores = model.evaluate(dataset, [metric], transformers)
		print(scores)

		# Predict the output and uncertainty.
		pred, std = model.predict_uncertainty(dataset)
		mean_error = np.mean(np.abs(dataset.y - pred))
		mean_value = np.mean(np.abs(dataset.y))
		mean_std = np.mean(std)
		print(mean_error, mean_value, mean_std)
		assert mean_error < 0.5 * mean_value
		assert mean_std > 0.5 * mean_error
		assert mean_std < mean_value
		@@ -121,7 +118,7 @@ class TestGraphModels(unittest.TestCase):

		featurizer = ConvMolFeaturizer(atom_properties=[atom_feature_name])
		X = featurizer.featurize(dataset.X)
		dataset = deepchem.data.NumpyDataset(X, np.array(y))
		dataset = dc.data.NumpyDataset(X, np.array(y))
		batch_size = 50
		model = GraphConvModel(
		len(tasks),
		@@ -173,3 +170,79 @@ class TestGraphModels(unittest.TestCase):
		module = model2.create_submodel(loss=loss)
		model2.restore()
		model2.fit(dataset, nb_epoch=1, submodel=module)

		def test_dag_model(self):
		tasks, dataset, transformers, metric = self.get_dataset(
		'classification', 'GraphConv')

		max_atoms = max([mol.get_num_atoms() for mol in dataset.X])
		transformer = dc.trans.DAGTransformer(max_atoms=max_atoms)
		dataset = transformer.transform(dataset)

		model = DAGModel(
		len(tasks), max_atoms=max_atoms, mode='classification', use_queue=False)

		model.fit(dataset, nb_epoch=10)
		scores = model.evaluate(dataset, [metric], transformers)
		assert scores['mean-roc_auc_score'] >= 0.9

		model.save()
		model = TensorGraph.load_from_dir(model.model_dir)
		scores2 = model.evaluate(dataset, [metric], transformers)
		assert np.allclose(scores['mean-roc_auc_score'],
		scores2['mean-roc_auc_score'])

		@attr("slow")
		def test_dag_regression_model(self):
		tasks, dataset, transformers, metric = self.get_dataset(
		'regression', 'GraphConv')

		max_atoms = max([mol.get_num_atoms() for mol in dataset.X])
		transformer = dc.trans.DAGTransformer(max_atoms=max_atoms)
		dataset = transformer.transform(dataset)

		model = DAGModel(
		len(tasks),
		max_atoms=max_atoms,
		mode='regression',
		learning_rate=0.003,
		use_queue=False)

		model.fit(dataset, nb_epoch=100)
		scores = model.evaluate(dataset, [metric], transformers)
		assert all(s < 0.15 for s in scores['mean_absolute_error'])

		model.save()
		model = TensorGraph.load_from_dir(model.model_dir)
		scores2 = model.evaluate(dataset, [metric], transformers)
		assert np.allclose(scores['mean_absolute_error'],
		scores2['mean_absolute_error'])

		@attr("slow")
		def test_dag_regression_uncertainty(self):
		tasks, dataset, transformers, metric = self.get_dataset(
		'regression', 'GraphConv')

		max_atoms = max([mol.get_num_atoms() for mol in dataset.X])
		transformer = dc.trans.DAGTransformer(max_atoms=max_atoms)
		dataset = transformer.transform(dataset)

		model = DAGModel(
		len(tasks),
		max_atoms=max_atoms,
		mode='regression',
		learning_rate=0.002,
		use_queue=False,
		dropout=0.1,
		uncertainty=True)

		model.fit(dataset, nb_epoch=100)

		# Predict the output and uncertainty.
		pred, std = model.predict_uncertainty(dataset)
		mean_error = np.mean(np.abs(dataset.y - pred))
		mean_value = np.mean(np.abs(dataset.y))
		mean_std = np.mean(std)
		assert mean_error < 0.5 * mean_value
		assert mean_std > 0.5 * mean_error
		assert mean_std < mean_value

deepchem/models/tensorgraph/tensor_graph.py

+6 −0

Original line number	Diff line number	Diff line
		@@ -1106,6 +1106,7 @@ class TensorGraph(Model):
		for key, value in d.items():
		if isinstance(key, Input):
		# Add or remove dimensions of size 1 to match the shape of the layer.
		try:
		value_dims = len(value.shape)
		layer_dims = len(key.shape)
		if value_dims < layer_dims:
		@@ -1116,6 +1117,8 @@ class TensorGraph(Model):
		if value_dims > layer_dims:
		if all(i == 1 for i in value.shape[layer_dims:]):
		value = tf.reshape(value, value.shape[:layer_dims])
		except:
		pass
		feed_dict[key] = tf.cast(value, key.dtype)
		else:
		feed_dict[key] = value
		@@ -1330,6 +1333,7 @@ def _enqueue_batch(tg, generator, graph, sess, n_enqueued, final_sample):
		if layer in feed_dict:
		value = feed_dict[layer]
		# Add or remove dimensions of size 1 to match the shape of the layer.
		try:
		value_dims = len(value.shape)
		layer_dims = len(layer.shape)
		if value_dims < layer_dims:
		@@ -1339,6 +1343,8 @@ def _enqueue_batch(tg, generator, graph, sess, n_enqueued, final_sample):
		if value_dims > layer_dims:
		if all(i == 1 for i in value.shape[layer_dims:]):
		value = value.reshape(value.shape[:layer_dims])
		except:
		pass
		else:
		value = np.zeros(
		[0] + list(layer.shape[1:]), dtype=layer.dtype.as_numpy_dtype)

deepchem/models/tensorgraph/graph_layers.py

+5 −5

File changed.

Contains only whitespace changes.

Admin message