:sparkles: create CGCNNModel (2e575fbb) · Commits · 钟慕尧 / deepchem

deepchem/models/init.py

+1 −1

Original line number	Diff line number	Diff line
		@@ -30,7 +30,7 @@ from deepchem.models.chemnet_models import Smiles2Vec, ChemCeption
		# PyTorch models
		try:
		from deepchem.models.torch_models import TorchModel
		from deepchem.models.torch_models import CGCNN, create_cgcnn_batch
		from deepchem.models.torch_models import CGCNN, CGCNNModel
		except ModuleNotFoundError:
		pass

deepchem/models/tests/test_cgcnn.py

+4 −7

Original line number	Diff line number	Diff line
		@@ -4,7 +4,7 @@ from os import path, remove
		from deepchem.feat import CGCNNFeaturizer
		from deepchem.molnet import load_perovskite
		from deepchem.metrics import Metric, mae_score
		from deepchem.models import TorchModel, CGCNN, create_cgcnn_batch, losses
		from deepchem.models import CGCNNModel, losses

		try:
		import dgl # noqa
		@@ -30,19 +30,16 @@ def test_cgcnn():
		train, valid, test = datasets

		# initialize models
		cgcnn = CGCNN(
		model = CGCNNModel(
		in_node_dim=92,
		hidden_node_dim=64,
		in_edge_dim=41,
		num_conv=3,
		predicator_hidden_feats=128,
		n_out=1)
		model = TorchModel(
		model=cgcnn,
		n_out=1,
		loss=losses.L2Loss(),
		batch_size=32,
		learning_rate=0.001,
		create_custom_batch=create_cgcnn_batch)
		learning_rate=0.001)

		# train
		model.fit(train, nb_epoch=50)

deepchem/models/torch_models/init.py

+1 −1

Original line number	Diff line number	Diff line
		# flake8:noqa
		from deepchem.models.torch_models.torch_model import TorchModel
		from deepchem.models.torch_models.cgcnn import CGCNN, create_cgcnn_batch
		from deepchem.models.torch_models.cgcnn import CGCNN, CGCNNModel

deepchem/models/torch_models/cgcnn.py

+98 −59

Original line number	Diff line number	Diff line
		@@ -3,50 +3,7 @@ import numpy as np
		import torch.nn as nn
		import torch.nn.functional as F


		def create_cgcnn_batch(batch, device):
		"""Create batch data for CGCNN.

		Parameters
		----------
		batch: Tuple
		The tuple are `(inputs, labels, weights)`.
		device: torch.device
		The device on which to run computations. If None, a device is
		chosen automatically.

		Returns
		-------
		inputs: DGLGraph
		DGLGraph for a batch of graphs.
		labels: List[torch.Tensor] or None
		The labels converted to torch.Tensor
		weights: List[torch.Tensor] or None
		The weights for each sample or sample/task pair converted to torch.Tensor

		Notes
		-----
		This class requires DGL and PyTorch to be installed.
		"""
		try:
		import dgl
		except:
		raise ValueError("This class requires DGL to be installed.")

		inputs, labels, weights = batch
		inputs = dgl.batch([graph.to_dgl_graph() for graph in inputs[0]]).to(device)
		if labels is not None:
		labels = [
		x.astype(np.float32) if x.dtype == np.float64 else x for x in labels
		]
		labels = [torch.as_tensor(x, device=device) for x in labels]
		if weights is not None:
		weights = [
		x.astype(np.float32) if x.dtype == np.float64 else x for x in weights
		]
		weights = [torch.as_tensor(x, device=device) for x in weights]

		return inputs, labels, weights
		from deepchem.models.torch_models.torch_model import TorchModel


		class CGCNNLayer(nn.Module):
		@@ -56,6 +13,7 @@ class CGCNNLayer(nn.Module):
		Please confirm how to use DGLGraph methods from below link.
		See: https://docs.dgl.ai/en/0.4.x/tutorials/models/1_gnn/9_gat.html
		"""

		def __init__(self,
		hidden_node_dim: int,
		edge_dim: int,
		@@ -111,19 +69,6 @@ class CGCNNLayer(nn.Module):
		class CGCNN(nn.Module):
		"""Crystal Graph Convolutional Neural Network (CGCNN).

		This class implements Crystal Graph Convolutional Neural Network (CGCNN).
		Here is a simple example of code that uses the CGCNN model with TorchModel
		and materials dataset.

		>> import deepchem as dc
		>> dataset_config = {"reload": False, "featurizer": dc.feat.CGCNNFeaturizer, "transformers": []}
		>> tasks, datasets, transformers = dc.molnet.load_perovskite(reload=False)
		>> train, valid, test = datasets
		>> cgcnn = dc.models.CGCNN()
		>> model = dc.models.TorchModel(cgcnn, loss=dc.models.losses.L2Loss(),
		.. create_custom_batch=dc.models.create_cgcnn_batch)
		>> model.fit(train, nb_epoch=50)

		This model takes arbitary crystal structures as an input, and predict material properties
		using the element information and connection of atoms in the crystal. If you want to get
		some material properties which has a high computational cost like band gap in the case
		@@ -159,11 +104,13 @@ class CGCNN(nn.Module):
		Parameters
		----------
		in_node_dim: int, default 92
		The length of the initial node feature vectors.
		The length of the initial node feature vectors. The 92 is
		based on length of vectors in the atom_init.json.
		hidden_node_dim: int, default 64
		The length of the hidden node feature vectors.
		in_edge_dim: int, default 41
		The length of the initial edge feature vectors.
		The length of the initial edge feature vectors. The 41 is
		based on default setting of CGCNNFeaturizer.
		num_conv: int, default 3
		The number of convolutional layers.
		predicator_hidden_feats: int, default 128
		@@ -214,3 +161,95 @@ class CGCNN(nn.Module):
		graph_feat = self.fc(graph_feat)
		out = self.out(graph_feat)
		return out


		class CGCNNModel(TorchModel):
		"""CGCNN wrapper model for converting PyTorch style to Keras style.

		Please confirm the details about CGCNN from CGCNN class docstring.
		Here is a simple example of code that uses the CGCNNModel with
		materials dataset.

		>> import deepchem as dc
		>> dataset_config = {"reload": False, "featurizer": dc.feat.CGCNNFeaturizer, "transformers": []}
		>> tasks, datasets, transformers = dc.molnet.load_perovskite(reload=False)
		>> train, valid, test = datasets
		>> model = dc.models.CGCNNModel(loss=dc.models.losses.L2Loss(), batch_size=32, learning_rate=0.001)
		>> model.fit(train, nb_epoch=50)

		Notes
		-----
		This class requires DGL and PyTorch to be installed.
		"""

		def __init__(self,
		in_node_dim: int = 92,
		hidden_node_dim: int = 64,
		in_edge_dim: int = 41,
		num_conv: int = 3,
		predicator_hidden_feats: int = 128,
		n_out: int = 1,
		**kwargs):
		"""
		Parameters
		----------
		in_node_dim: int, default 92
		The length of the initial node feature vectors. The 92 is
		based on length of vectors in the atom_init.json.
		hidden_node_dim: int, default 64
		The length of the hidden node feature vectors.
		in_edge_dim: int, default 41
		The length of the initial edge feature vectors. The 41 is
		based on default setting of CGCNNFeaturizer.
		num_conv: int, default 3
		The number of convolutional layers.
		predicator_hidden_feats: int, default 128
		Size of hidden graph representations in the predicator, default to 128.
		n_out: int
		Number of the output size, default to 1.
		"""
		model = CGCNN(in_node_dim, hidden_node_dim, in_edge_dim, num_conv,
		predicator_hidden_feats, n_out)
		super(CGCNNModel, self).__init__(model, **kwargs)

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

		Parameters
		----------
		batch: Tuple
		The tuple are `(inputs, labels, weights)`.

		Returns
		-------
		inputs: DGLGraph
		DGLGraph for a batch of graphs.
		labels: List[torch.Tensor] or None
		The labels converted to torch.Tensor
		weights: List[torch.Tensor] or None
		The weights for each sample or sample/task pair converted to torch.Tensor

		Notes
		-----
		This class requires DGL and PyTorch to be installed.
		"""
		try:
		import dgl
		except:
		raise ValueError("This class requires DGL to be installed.")

		inputs, labels, weights = batch
		inputs = dgl.batch([graph.to_dgl_graph() for graph in inputs[0]]).to(
		self.device)
		if labels is not None:
		labels = [
		x.astype(np.float32) if x.dtype == np.float64 else x for x in labels
		]
		labels = [torch.as_tensor(x, device=self.device) for x in labels]
		if weights is not None:
		weights = [
		x.astype(np.float32) if x.dtype == np.float64 else x for x in weights
		]
		weights = [torch.as_tensor(x, device=self.device) for x in weights]

		return inputs, labels, weights

deepchem/models/torch_models/torch_model.py

+0 −8

Original line number	Diff line number	Diff line
		@@ -122,7 +122,6 @@ class TorchModel(Model):
		wandb: bool = False,
		log_frequency: int = 100,
		device: Optional[torch.device] = None,
		create_custom_batch: Callable[[Tuple, torch.device], Tuple] = None,
		**kwargs) -> None:
		"""Create a new TorchModel.

		@@ -161,9 +160,6 @@ class TorchModel(Model):
		device: torch.device
		the device on which to run computations. If None, a device is
		chosen automatically.
		create_custom_batch: function, default None
		This function makes user-defined batch data. This function takes two arguments,
		`batch`, `device` and returns user-defined batch data.
		"""
		super(TorchModel, self).__init__(
		model_instance=model, model_dir=model_dir, **kwargs)
		@@ -188,7 +184,6 @@ class TorchModel(Model):
		device = torch.device('cpu')
		self.device = device
		self.model.to(device)
		self.create_custom_batch = create_custom_batch

		# W&B logging
		if wandb and not is_wandb_available():
		@@ -844,9 +839,6 @@ class TorchModel(Model):

		def _prepare_batch(self,
		batch: Tuple[Any, Any, Any]) -> Tuple[List, List, List]:
		if self.create_custom_batch is not None:
		return self.create_custom_batch(batch, self.device)

		inputs, labels, weights = batch
		inputs = [
		x.astype(np.float32) if x.dtype == np.float64 else x for x in inputs

Admin message