Update (93898573) · Commits · 钟慕尧 / deepchem

deepchem/models/init.py

+1 −0

Original line number	Diff line number	Diff line
		@@ -31,6 +31,7 @@ from deepchem.models.gbdt_models import GBDTModel
		# PyTorch models
		try:
		from deepchem.models.torch_models import TorchModel
		from deepchem.models.torch_models import AttentiveFP, AttentiveFPModel
		from deepchem.models.torch_models import CGCNN, CGCNNModel
		from deepchem.models.torch_models import GAT, GATModel
		from deepchem.models.torch_models import GCN, GCNModel

deepchem/models/torch_models/init.py

+1 −0

Original line number	Diff line number	Diff line
		# flake8:noqa
		from deepchem.models.torch_models.torch_model import TorchModel
		from deepchem.models.torch_models.attentivefp import AttentiveFP, AttentiveFPModel
		from deepchem.models.torch_models.cgcnn import CGCNN, CGCNNModel
		from deepchem.models.torch_models.gat import GAT, GATModel
		from deepchem.models.torch_models.gcn import GCN, GCNModel

deepchem/models/torch_models/attentivefp.py

0 → 100644

+321 −0

Original line number	Diff line number	Diff line
		"""
		DGL-based AttentiveFP for graph property prediction.
		"""
		import torch.nn as nn
		import torch.nn.functional as F

		from deepchem.models.losses import Loss, L2Loss, SparseSoftmaxCrossEntropy
		from deepchem.models.torch_models.torch_model import TorchModel

		class AttentiveFP(nn.Module):
		"""Model for Graph Property Prediction.

		This model proceeds as follows:

		* Combine node features and edge features for initializing node representations,
		which involves a round of message passing
		* Update node representations with multiple rounds of message passing
		* For each graph, compute its representation by combining the representations
		of all nodes in it, which involves a gated recurrent unit (GRU).
		* Perform the final prediction using a linear layer

		Examples
		--------

		>>> import deepchem as dc
		>>> import dgl
		>>> from deepchem.models import AttentiveFP
		>>> smiles = ["C1CCC1", "C1=CC=CN=C1"]
		>>> featurizer = dc.feat.MolGraphConvFeaturizer(use_edges=True)
		>>> graphs = featurizer.featurize(smiles)
		>>> print(type(graphs[0]))
		<class 'deepchem.feat.graph_data.GraphData'>
		>>> dgl_graphs = [graphs[i].to_dgl_graph() for i in range(len(graphs))]
		>>> # Batch two graphs into a graph of two connected components
		>>> batch_dgl_graph = dgl.batch(dgl_graphs)
		>>> model = AttentiveFP(n_tasks=1, mode='regression')
		>>> preds = model(batch_dgl_graph)
		>>> print(type(preds))
		<class 'torch.Tensor'>
		>>> preds.shape == (2, 1)
		True

		References
		----------
		.. [1] Zhaoping Xiong, Dingyan Wang, Xiaohong Liu, Feisheng Zhong, Xiaozhe Wan, Xutong Li,
		Zhaojun Li, Xiaomin Luo, Kaixian Chen, Hualiang Jiang, and Mingyue Zheng. "Pushing
		the Boundaries of Molecular Representation for Drug Discovery with the Graph Attention
		Mechanism." Journal of Medicinal Chemistry. 2020, 63, 16, 8749–8760.

		Notes
		-----
		This class requires DGL (https://github.com/dmlc/dgl) and DGL-LifeSci
		(https://github.com/awslabs/dgl-lifesci) to be installed.
		"""

		def __init__(self,
		n_tasks: int,
		num_layers: int = 2,
		num_timesteps: int = 2,
		graph_feat_size: int = 200,
		dropout: float = 0.,
		mode: str = 'regression',
		number_atom_features: int = 30,
		number_bond_features: int = 11,
		n_classes: int = 2,
		nfeat_name: str = 'x',
		efeat_name: str = 'edge_attr'):
		"""
		Parameters
		----------
		n_tasks: int
		Number of tasks.
		num_layers: int
		Number of graph neural network layers, i.e. number of rounds of message passing.
		Default to 2.
		num_timesteps: int
		Number of time steps for updating graph representations with a GRU. Default to 2.
		graph_feat_size: int
		Size for graph representations. Default to 200.
		dropout: float
		Dropout probability. Default to 0.
		mode: str
		The model type, 'classification' or 'regression'. Default to 'regression'.
		number_atom_features: int
		The length of the initial atom feature vectors. Default to 30.
		number_bond_features: int
		The length of the initial bond feature vectors. Default to 11.
		n_classes: int
		The number of classes to predict per task
		(only used when ``mode`` is 'classification'). Default to 2.
		nfeat_name: str
		For an input graph ``g``, the model assumes that it stores node features in
		``g.ndata[nfeat_name]`` and will retrieve input node features from that.
		Default to 'x'.
		efeat_name: str
		For an input graph ``g``, the model assumes that it stores edge features in
		``g.edata[efeat_name]`` and will retrieve input edge features from that.
		Default to 'edge_attr'.
		"""
		try:
		import dgl
		except:
		raise ImportError('This class requires dgl.')
		try:
		import dgllife
		except:
		raise ImportError('This class requires dgllife.')

		if mode not in ['classification', 'regression']:
		raise ValueError("mode must be either 'classification' or 'regression'")

		super(AttentiveFP, self).__init__()

		self.n_tasks = n_tasks
		self.mode = mode
		self.n_classes = n_classes
		self.nfeat_name = nfeat_name
		self.efeat_name = efeat_name
		if mode == 'classification':
		out_size = n_tasks * n_classes
		else:
		out_size = n_tasks

		from dgllife.model import AttentiveFPPredictor as DGLAttentiveFPPredictor

		self.model = DGLAttentiveFPPredictor(node_feat_size=number_atom_features,
		edge_feat_size=number_bond_features,
		num_layers=num_layers,
		num_timesteps=num_timesteps,
		graph_feat_size=graph_feat_size,
		n_tasks=out_size,
		dropout=dropout)

		def forward(self, g):
		"""Predict graph labels

		Parameters
		----------
		g: DGLGraph
		A DGLGraph for a batch of graphs. It stores the node features in
		``dgl_graph.ndata[self.nfeat_name]`` and edge features in
		``dgl_graph.edata[self.efeat_name]``.

		Returns
		-------
		torch.Tensor
		The model output.

		* When self.mode = 'regression',
		its shape will be ``(dgl_graph.batch_size, self.n_tasks)``.
		* When self.mode = 'classification', the output consists of probabilities
		for classes. Its shape will be
		``(dgl_graph.batch_size, self.n_tasks, self.n_classes)`` if self.n_tasks > 1;
		its shape will be ``(dgl_graph.batch_size, self.n_classes)`` if self.n_tasks is 1.
		torch.Tensor, optional
		This is only returned when self.mode = 'classification', the output consists of the
		logits for classes before softmax.
		"""
		node_feats = g.ndata[self.nfeat_name]
		edge_feats = g.edata[self.efeat_name]
		out = self.model(g, node_feats, edge_feats)

		if self.mode == 'classification':
		if self.n_tasks == 1:
		logits = out.view(-1, self.n_classes)
		softmax_dim = 1
		else:
		logits = out.view(-1, self.n_tasks, self.n_classes)
		softmax_dim = 2
		proba = F.softmax(logits, dim=softmax_dim)
		return proba, logits
		else:
		return out


		class AttentiveFPModel(TorchModel):
		"""Model for Graph Property Prediction.

		This model proceeds as follows:

		* Combine node features and edge features for initializing node representations,
		which involves a round of message passing
		* Update node representations with multiple rounds of message passing
		* For each graph, compute its representation by combining the representations
		of all nodes in it, which involves a gated recurrent unit (GRU).
		* Perform the final prediction using a linear layer

		Examples
		--------

		>>>
		>> import deepchem as dc
		>> from deepchem.models import AttentiveFPModel
		>> featurizer = dc.feat.MolGraphConvFeaturizer(use_edges=True)
		>> tasks, datasets, transformers = dc.molnet.load_tox21(
		.. reload=False, featurizer=featurizer, transformers=[])
		>> train, valid, test = datasets
		>> model = dc.models.AttentiveFPModel(mode='classification', n_tasks=len(tasks),
		.. batch_size=32, learning_rate=0.001)
		>> model.fit(train, nb_epoch=50)

		References
		----------
		.. [1] Zhaoping Xiong, Dingyan Wang, Xiaohong Liu, Feisheng Zhong, Xiaozhe Wan, Xutong Li,
		Zhaojun Li, Xiaomin Luo, Kaixian Chen, Hualiang Jiang, and Mingyue Zheng. "Pushing
		the Boundaries of Molecular Representation for Drug Discovery with the Graph
		Attention Mechanism." Journal of Medicinal Chemistry. 2020, 63, 16, 8749–8760.

		Notes
		-----
		This class requires DGL (https://github.com/dmlc/dgl) and DGL-LifeSci
		(https://github.com/awslabs/dgl-lifesci) to be installed.
		"""

		def __init__(self,
		n_tasks: int,
		num_layers: int = 2,
		num_timesteps: int = 2,
		graph_feat_size: int = 200,
		dropout: float = 0.,
		mode: str = 'regression',
		number_atom_features: int = 30,
		number_bond_features: int = 11,
		n_classes: int = 2,
		nfeat_name: str = 'x',
		efeat_name: str = 'edge_attr',
		self_loop: bool = True,
		**kwargs):
		"""
		Parameters
		----------
		n_tasks: int
		Number of tasks.
		num_layers: int
		Number of graph neural network layers, i.e. number of rounds of message passing.
		Default to 2.
		num_timesteps: int
		Number of time steps for updating graph representations with a GRU. Default to 2.
		graph_feat_size: int
		Size for graph representations. Default to 200.
		dropout: float
		Dropout probability. Default to 0.
		mode: str
		The model type, 'classification' or 'regression'. Default to 'regression'.
		number_atom_features: int
		The length of the initial atom feature vectors. Default to 30.
		number_bond_features: int
		The length of the initial bond feature vectors. Default to 11.
		n_classes: int
		The number of classes to predict per task
		(only used when ``mode`` is 'classification'). Default to 2.
		nfeat_name: str
		For an input graph ``g``, the model assumes that it stores node features in
		``g.ndata[nfeat_name]`` and will retrieve input node features from that.
		Default to 'x'.
		efeat_name: str
		For an input graph ``g``, the model assumes that it stores edge features in
		``g.edata[efeat_name]`` and will retrieve input edge features from that.
		Default to 'edge_attr'.
		self_loop: bool
		Whether to add self loops for the nodes, i.e. edges from nodes to themselves.
		Default to True.
		kwargs
		This can include any keyword argument of TorchModel.
		"""
		model = AttentiveFP(
		n_tasks=n_tasks,
		num_layers=num_layers,
		num_timesteps=num_timesteps,
		graph_feat_size=graph_feat_size,
		dropout=dropout,
		mode=mode,
		number_atom_features=number_atom_features,
		number_bond_features=number_bond_features,
		n_classes=n_classes,
		nfeat_name=nfeat_name,
		efeat_name=efeat_name)
		if mode == 'regression':
		loss: Loss = L2Loss()
		output_types = ['prediction']
		else:
		loss = SparseSoftmaxCrossEntropy()
		output_types = ['prediction', 'loss']
		super(AttentiveFPModel, self).__init__(
		model, loss=loss, output_types=output_types, **kwargs)

		self._self_loop = self_loop

		def _prepare_batch(self, batch):
		"""Create batch data for AttentiveFP.

		Parameters
		----------
		batch: tuple
		The tuple is ``(inputs, labels, weights)``.
		self_loop: bool
		Whether to add self loops for the nodes, i.e. edges from nodes
		to themselves. Default to False.

		Returns
		-------
		inputs: DGLGraph
		DGLGraph for a batch of graphs.
		labels: list of torch.Tensor or None
		The graph labels.
		weights: list of torch.Tensor or None
		The weights for each sample or sample/task pair converted to torch.Tensor.
		"""
		try:
		import dgl
		except:
		raise ImportError('This class requires dgl.')

		inputs, labels, weights = batch
		dgl_graphs = [
		graph.to_dgl_graph(self_loop=self._self_loop) for graph in inputs[0]
		]
		inputs = dgl.batch(dgl_graphs).to(self.device)
		_, labels, weights = super(AttentiveFPModel, self)._prepare_batch(([], labels,
		weights))
		return inputs, labels, weights

deepchem/models/torch_models/gat.py

+12 −10

Original line number	Diff line number	Diff line
		@@ -34,7 +34,7 @@ class GAT(nn.Module):
		>>> dgl_graphs = [graphs[i].to_dgl_graph() for i in range(len(graphs))]
		>>> # Batch two graphs into a graph of two connected components
		>>> batch_dgl_graph = dgl.batch(dgl_graphs)
		>>> model = GAT(n_tasks=1, number_atom_features=30, mode='regression')
		>>> model = GAT(n_tasks=1, mode='regression')
		>>> preds = model(batch_dgl_graph)
		>>> print(type(preds))
		<class 'torch.Tensor'>
		@@ -64,7 +64,7 @@ class GAT(nn.Module):
		predictor_hidden_feats: int = 128,
		predictor_dropout: float = 0.,
		mode: str = 'regression',
		number_atom_features: int = 75,
		number_atom_features: int = 30,
		n_classes: int = 2,
		nfeat_name: str = 'x'):
		"""
		@@ -101,15 +101,16 @@ class GAT(nn.Module):
		predictor_dropout: float
		The dropout probability in the output MLP predictor. Default to 0.
		mode: str
		The model type, 'classification' or 'regression'.
		The model type, 'classification' or 'regression'. Default to 'regression'.
		number_atom_features: int
		The length of the initial atom feature vectors. Default to 75.
		The length of the initial atom feature vectors. Default to 30.
		n_classes: int
		The number of classes to predict per task
		(only used when ``mode`` is 'classification').
		(only used when ``mode`` is 'classification'). Default to 2.
		nfeat_name: str
		For an input graph ``g``, the model assumes that it stores node features in
		``g.ndata[nfeat_name]`` and will retrieve input node features from that.
		Default to 'x'.
		"""
		try:
		import dgl
		@@ -235,7 +236,7 @@ class GATModel(TorchModel):
		.. reload=False, featurizer=featurizer, transformers=[])
		>> train, valid, test = datasets
		>> model = dc.models.GATModel(mode='classification', n_tasks=len(tasks),
		.. number_atom_features=30, batch_size=32, learning_rate=0.001)
		.. batch_size=32, learning_rate=0.001)
		>> model.fit(train, nb_epoch=50)

		References
		@@ -261,7 +262,7 @@ class GATModel(TorchModel):
		predictor_hidden_feats: int = 128,
		predictor_dropout: float = 0.,
		mode: str = 'regression',
		number_atom_features: int = 75,
		number_atom_features: int = 30,
		n_classes: int = 2,
		nfeat_name: str = 'x',
		self_loop: bool = True,
		@@ -300,15 +301,16 @@ class GATModel(TorchModel):
		predictor_dropout: float
		The dropout probability in the output MLP predictor. Default to 0.
		mode: str
		The model type, 'classification' or 'regression'.
		The model type, 'classification' or 'regression'. Default to 'regression'.
		number_atom_features: int
		The length of the initial atom feature vectors. Default to 75.
		The length of the initial atom feature vectors. Default to 30.
		n_classes: int
		The number of classes to predict per task
		(only used when ``mode`` is 'classification').
		(only used when ``mode`` is 'classification'). Default to 2.
		nfeat_name: str
		For an input graph ``g``, the model assumes that it stores node features in
		``g.ndata[nfeat_name]`` and will retrieve input node features from that.
		Default to 'x'.
		self_loop: bool
		Whether to add self loops for the nodes, i.e. edges from nodes to themselves.
		Default to True.

deepchem/models/torch_models/gcn.py

+12 −10

Original line number	Diff line number	Diff line
		@@ -34,7 +34,7 @@ class GCN(nn.Module):
		>>> dgl_graphs = [graphs[i].to_dgl_graph() for i in range(len(graphs))]
		>>> # Batch two graphs into a graph of two connected components
		>>> batch_dgl_graph = dgl.batch(dgl_graphs)
		>>> model = GCN(n_tasks=1, number_atom_features=30, mode='regression')
		>>> model = GCN(n_tasks=1, mode='regression')
		>>> preds = model(batch_dgl_graph)
		>>> print(type(preds))
		<class 'torch.Tensor'>
		@@ -77,7 +77,7 @@ class GCN(nn.Module):
		predictor_hidden_feats: int = 128,
		predictor_dropout: float = 0.,
		mode: str = 'regression',
		number_atom_features: int = 75,
		number_atom_features: int = 30,
		n_classes: int = 2,
		nfeat_name: str = 'x'):
		"""
		@@ -103,15 +103,16 @@ class GCN(nn.Module):
		predictor_dropout: float
		The dropout probability in the output MLP predictor. Default to 0.
		mode: str
		The model type, 'classification' or 'regression'.
		The model type, 'classification' or 'regression'. Default to 'regression'.
		number_atom_features: int
		The length of the initial atom feature vectors. Default to 75.
		The length of the initial atom feature vectors. Default to 30.
		n_classes: int
		The number of classes to predict per task
		(only used when ``mode`` is 'classification').
		(only used when ``mode`` is 'classification'). Default to 2.
		nfeat_name: str
		For an input graph ``g``, the model assumes that it stores node features in
		``g.ndata[nfeat_name]`` and will retrieve input node features from that.
		Default to 'x'.
		"""
		try:
		import dgl
		@@ -219,7 +220,7 @@ class GCNModel(TorchModel):
		.. reload=False, featurizer=featurizer, transformers=[])
		>> train, valid, test = datasets
		>> model = dc.models.GCNModel(mode='classification', n_tasks=len(tasks),
		.. number_atom_features=30, batch_size=32, learning_rate=0.001)
		.. batch_size=32, learning_rate=0.001)
		>> model.fit(train, nb_epoch=50)

		References
		@@ -258,7 +259,7 @@ class GCNModel(TorchModel):
		predictor_hidden_feats: int = 128,
		predictor_dropout: float = 0.,
		mode: str = 'regression',
		number_atom_features=75,
		number_atom_features=30,
		n_classes: int = 2,
		nfeat_name: str = 'x',
		self_loop: bool = True,
		@@ -286,15 +287,16 @@ class GCNModel(TorchModel):
		predictor_dropout: float
		The dropout probability in the output MLP predictor. Default to 0.
		mode: str
		The model type, 'classification' or 'regression'.
		The model type, 'classification' or 'regression'. Default to 'regression'.
		number_atom_features: int
		The length of the initial atom feature vectors. Default to 75.
		The length of the initial atom feature vectors. Default to 30.
		n_classes: int
		The number of classes to predict per task
		(only used when ``mode`` is 'classification').
		(only used when ``mode`` is 'classification'). Default to 2.
		nfeat_name: str
		For an input graph ``g``, the model assumes that it stores node features in
		``g.ndata[nfeat_name]`` and will retrieve input node features from that.
		Default to 'x'.
		self_loop: bool
		Whether to add self loops for the nodes, i.e. edges from nodes to themselves.
		Default to True.

Admin message