Implemented uncertainty for MPNNModel (6877fed6) · Commits · 钟慕尧 / deepchem

deepchem/models/tensorgraph/models/graph_models.py

+25 −4

Original line number	Diff line number	Diff line
		@@ -763,7 +763,9 @@ class MPNNModel(TensorGraph):
		T=5,
		M=10,
		mode="regression",
		dropout=0.0,
		n_classes=2,
		uncertainty=False,
		**kwargs):
		"""
		Parameters
		@@ -778,9 +780,13 @@ class MPNNModel(TensorGraph):
		Number of units(convolution depths) in corresponding hidden layer
		n_graph_feat: int, optional
		Number of output features for each molecule(graph)
		dropout: float
		the dropout probablity to use.
		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
		"""
		if mode not in ['classification', 'regression']:
		raise ValueError("mode must be either 'classification' or 'regression'")
		@@ -792,6 +798,12 @@ class MPNNModel(TensorGraph):
		self.M = M
		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(MPNNModel, self).__init__(**kwargs)
		self.build_graph()

		@@ -829,9 +841,7 @@ class MPNNModel(TensorGraph):
		labels = Label(shape=(None, n_tasks, n_classes))
		logits = Reshape(
		shape=(None, n_tasks, n_classes),
		in_layers=[
		Dense(in_layers=dense1, out_channels=n_tasks * n_classes)
		])
		in_layers=[Dense(in_layers=dense1, out_channels=n_tasks * n_classes)])
		logits = TrimGraphOutput([logits, weights])
		output = SoftMax(logits)
		self.add_output(output)
		@@ -845,6 +855,17 @@ class MPNNModel(TensorGraph):
		in_layers=[Dense(in_layers=dense1, out_channels=n_tasks)])
		output = TrimGraphOutput([output, weights])
		self.add_output(output)
		if self.uncertainty:
		log_var = Reshape(
		shape=(None, n_tasks),
		in_layers=[Dense(in_layers=dense1, out_channels=n_tasks)])
		log_var = TrimGraphOutput([log_var, weights])
		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)

deepchem/models/tensorgraph/models/test_graph_models.py

+43 −2

Original line number	Diff line number	Diff line
		@@ -284,7 +284,14 @@ class TestGraphModels(unittest.TestCase):
		tasks, dataset, transformers, metric = self.get_dataset(
		'classification', 'Weave')

		model = MPNNModel(len(tasks), mode='classification', n_hidden=75, n_atom_feat=75, n_pair_feat=14, T=1, M=1)
		model = MPNNModel(
		len(tasks),
		mode='classification',
		n_hidden=75,
		n_atom_feat=75,
		n_pair_feat=14,
		T=1,
		M=1)

		model.fit(dataset, nb_epoch=20)
		scores = model.evaluate(dataset, [metric], transformers)
		@@ -301,7 +308,14 @@ class TestGraphModels(unittest.TestCase):
		tasks, dataset, transformers, metric = self.get_dataset(
		'regression', 'Weave')

		model = MPNNModel(len(tasks), mode='regression', n_hidden=75, n_atom_feat=75, n_pair_feat=14, T=1, M=1)
		model = MPNNModel(
		len(tasks),
		mode='regression',
		n_hidden=75,
		n_atom_feat=75,
		n_pair_feat=14,
		T=1,
		M=1)

		model.fit(dataset, nb_epoch=50)
		scores = model.evaluate(dataset, [metric], transformers)
		@@ -312,3 +326,30 @@ class TestGraphModels(unittest.TestCase):
		scores2 = model.evaluate(dataset, [metric], transformers)
		assert np.allclose(scores['mean_absolute_error'],
		scores2['mean_absolute_error'])

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

		model = MPNNModel(
		len(tasks),
		mode='regression',
		n_hidden=75,
		n_atom_feat=75,
		n_pair_feat=14,
		T=1,
		M=1,
		dropout=0.1,
		uncertainty=True)

		model.fit(dataset, nb_epoch=40)

		# 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

Admin message