Skills › Scientific Agent Skills › deepchem › graph_neural_network.py
graph_neural_network.py
334 lines · skills/deepchem/scripts/graph_neural_network.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3Graph Neural Network Training Script
5This script demonstrates training Graph Convolutional Networks (GCNs) and other
6graph-based models for molecular property prediction.
8Usage:
9 python graph_neural_network.py --dataset tox21 --model gcn
10 python graph_neural_network.py --dataset bbbp --model attentivefp
11 python graph_neural_network.py --data custom.csv --task-type regression
12"""
14import argparse
15import deepchem as dc
16import sys
17from _common import load_csv, validate_splits, evaluate_splits
20AVAILABLE_MODELS = {
21 'gcn': 'Graph Convolutional Network',
22 'gat': 'Graph Attention Network',
23 'attentivefp': 'Attentive Fingerprint',
24 'mpnn': 'Message Passing Neural Network',
25 'dmpnn': 'Directed Message Passing Neural Network'
26}
28MOLNET_DATASETS = {
29 'tox21': ('classification', 12),
30 'bbbp': ('classification', 1),
31 'bace': ('classification', 1),
32 'hiv': ('classification', 1),
33 'delaney': ('regression', 1),
34 'freesolv': ('regression', 1),
35 'lipo': ('regression', 1)
36}
38# MoleculeNet loader names that are not simply load_<dataset>. BACE ships as
39# separate classification and regression loaders, so dc.molnet.load_bace does
40# not exist; the entry above treats it as a classification benchmark.
41MOLNET_LOADERS = {
42 'bace': 'load_bace_classification'
43}
46def molnet_loader(dataset_name):
47 """Return the dc.molnet loader function for a MoleculeNet dataset name."""
48 return getattr(dc.molnet, MOLNET_LOADERS.get(dataset_name, f'load_{dataset_name}'))
51def create_model(model_type, n_tasks, mode='classification'):
52 """
53 Create a graph neural network model.
55 Args:
56 model_type: Type of model ('gcn', 'gat', 'attentivefp', etc.)
57 n_tasks: Number of prediction tasks
58 mode: 'classification' or 'regression'
60 Returns:
61 DeepChem model
62 """
63 if model_type == 'gcn':
64 return dc.models.GCNModel(
65 n_tasks=n_tasks,
66 mode=mode,
67 device="cpu",
68 batch_size=128,
69 learning_rate=0.001,
70 dropout=0.0
71 )
72 elif model_type == 'gat':
73 return dc.models.GATModel(
74 n_tasks=n_tasks,
75 mode=mode,
76 device="cpu",
77 batch_size=128,
78 learning_rate=0.001
79 )
80 elif model_type == 'attentivefp':
81 return dc.models.AttentiveFPModel(
82 n_tasks=n_tasks,
83 mode=mode,
84 device="cpu",
85 batch_size=128,
86 learning_rate=0.001
87 )
88 elif model_type == 'mpnn':
89 from deepchem.models.torch_models import MPNNModel
90 return MPNNModel(
91 n_tasks=n_tasks,
92 mode=mode,
93 device="cpu",
94 batch_size=128,
95 learning_rate=0.001
96 )
97 elif model_type == 'dmpnn':
98 return dc.models.DMPNNModel(
99 n_classes=2,
100 n_tasks=n_tasks,
101 mode=mode,
102 device="cpu",
103 batch_size=128,
104 learning_rate=0.001
105 )
106 else:
107 raise ValueError(f"Unknown model type: {model_type}")
110def create_featurizer(model_type):
111 """Match the released model's representation, including required bonds."""
112 if model_type == 'dmpnn':
113 return dc.feat.DMPNNFeaturizer()
114 if model_type in AVAILABLE_MODELS:
115 return dc.feat.MolGraphConvFeaturizer(
116 use_edges=model_type in ('attentivefp', 'mpnn'))
117 raise ValueError(f"Unknown model type: {model_type}")
120def train_on_molnet(dataset_name, model_type, n_epochs=50):
121 """
122 Train a graph neural network on a MoleculeNet benchmark dataset.
124 Args:
125 dataset_name: Name of MoleculeNet dataset
126 model_type: Type of model to train
127 n_epochs: Number of training epochs
129 Returns:
130 Trained model and test scores
131 """
132 print("=" * 70)
133 print(f"Training {AVAILABLE_MODELS[model_type]} on {dataset_name.upper()}")
134 print("=" * 70)
136 # Get dataset info
137 task_type, n_tasks_default = MOLNET_DATASETS[dataset_name]
139 # Load dataset with graph featurization
140 print(f"\nLoading {dataset_name} dataset with {model_type} features...")
141 load_func = molnet_loader(dataset_name)
142 tasks, datasets, transformers = load_func(
143 featurizer=create_featurizer(model_type),
144 splitter='scaffold'
145 )
146 train, valid, test = datasets
148 n_tasks = len(tasks)
149 print(f"\nDataset Information:")
150 print(f" Task type: {task_type}")
151 print(f" Number of tasks: {n_tasks}")
152 print(f" Training samples: {len(train)}")
153 print(f" Validation samples: {len(valid)}")
154 print(f" Test samples: {len(test)}")
156 validate_splits((train, valid, test), task_type, len(tasks))
158 # Create model
159 print(f"\nCreating {AVAILABLE_MODELS[model_type]} model...")
160 model = create_model(model_type, n_tasks, mode=task_type)
162 # Train
163 print(f"\nTraining for {n_epochs} epochs...")
164 model.fit(train, nb_epoch=n_epochs)
165 print("Training complete!")
167 # Evaluate
168 print("\n" + "=" * 70)
169 print("Model Evaluation")
170 print("=" * 70)
172 results = evaluate_splits(model, (train, valid, test), task_type, transformers)
174 return model, results
177def train_on_custom_data(data_path, model_type, task_type, target_cols, smiles_col='smiles', n_epochs=50):
178 """
179 Train a graph neural network on custom CSV data.
181 Args:
182 data_path: Path to CSV file
183 model_type: Type of model to train
184 task_type: 'classification' or 'regression'
185 target_cols: List of target column names
186 smiles_col: Name of SMILES column
187 n_epochs: Number of training epochs
189 Returns:
190 Trained model and test dataset
191 """
192 print("=" * 70)
193 print(f"Training {AVAILABLE_MODELS[model_type]} on Custom Data")
194 print("=" * 70)
196 # Load and featurize data
197 print(f"\nLoading data from {data_path}...")
198 featurizer = create_featurizer(model_type)
199 dataset = load_csv(data_path, target_cols, smiles_col, featurizer)
201 print(f"Loaded {len(dataset)} molecules")
203 # Split data
204 print("\nSplitting data with scaffold splitter...")
205 splitter = dc.splits.ScaffoldSplitter()
206 train, valid, test = splitter.train_valid_test_split(
207 dataset,
208 frac_train=0.8,
209 frac_valid=0.1,
210 frac_test=0.1
211 )
213 print(f" Training: {len(train)}")
214 print(f" Validation: {len(valid)}")
215 print(f" Test: {len(test)}")
217 # Create model
218 print(f"\nCreating {AVAILABLE_MODELS[model_type]} model...")
219 n_tasks = len(target_cols)
220 validate_splits((train, valid, test), task_type, n_tasks)
221 model = create_model(model_type, n_tasks, mode=task_type)
223 # Train
224 print(f"\nTraining for {n_epochs} epochs...")
225 model.fit(train, nb_epoch=n_epochs)
226 print("Training complete!")
228 # Evaluate
229 print("\n" + "=" * 70)
230 print("Model Evaluation")
231 print("=" * 70)
233 evaluate_splits(model, (train, valid, test), task_type)
235 return model, test
238def main():
239 parser = argparse.ArgumentParser(
240 description='Train graph neural networks for molecular property prediction'
241 )
242 parser.add_argument(
243 '--model',
244 type=str,
245 choices=list(AVAILABLE_MODELS.keys()),
246 default='gcn',
247 help='Type of graph neural network model'
248 )
249 parser.add_argument(
250 '--dataset',
251 type=str,
252 choices=list(MOLNET_DATASETS.keys()),
253 default=None,
254 help='MoleculeNet dataset to use'
255 )
256 parser.add_argument(
257 '--data',
258 type=str,
259 default=None,
260 help='Path to custom CSV file'
261 )
262 parser.add_argument(
263 '--task-type',
264 type=str,
265 choices=['classification', 'regression'],
266 default='classification',
267 help='Type of prediction task (for custom data)'
268 )
269 parser.add_argument(
270 '--targets',
271 nargs='+',
272 default=['target'],
273 help='Names of target columns (for custom data)'
274 )
275 parser.add_argument(
276 '--smiles-col',
277 type=str,
278 default='smiles',
279 help='Name of SMILES column'
280 )
281 parser.add_argument(
282 '--epochs',
283 type=int,
284 default=50,
285 help='Number of training epochs'
286 )
288 args = parser.parse_args()
289 if args.epochs < 1:
290 parser.error('--epochs must be positive')
292 # Validate arguments
293 if args.dataset is None and args.data is None:
294 print("Error: Must specify either --dataset (MoleculeNet) or --data (custom CSV)",
295 file=sys.stderr)
296 return 1
298 if args.dataset and args.data:
299 print("Error: Cannot specify both --dataset and --data",
300 file=sys.stderr)
301 return 1
303 # Train model
304 try:
305 if args.dataset:
306 model, results = train_on_molnet(
307 args.dataset,
308 args.model,
309 n_epochs=args.epochs
310 )
311 else:
312 model, test_set = train_on_custom_data(
313 args.data,
314 args.model,
315 args.task_type,
316 args.targets,
317 smiles_col=args.smiles_col,
318 n_epochs=args.epochs
319 )
321 print("\n" + "=" * 70)
322 print("Training Complete!")
323 print("=" * 70)
324 return 0
326 except Exception as e:
327 print(f"\nError: {e}", file=sys.stderr)
328 import traceback
329 traceback.print_exc()
330 return 1
333if __name__ == '__main__':
334 sys.exit(main())