Initial commit
This commit is contained in:
Zhongwei Li
529
skills/torch_geometric/scripts/create_gnn_template.py
Normal file
529
skills/torch_geometric/scripts/create_gnn_template.py
Normal file
@@ -0,0 +1,529 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Generate boilerplate code for common GNN architectures in PyTorch Geometric.
|
||||
|
||||
This script creates ready-to-use GNN model templates with training loops,
|
||||
evaluation metrics, and proper data handling.
|
||||
|
||||
Usage:
|
||||
python create_gnn_template.py --model gcn --task node_classification --output my_model.py
|
||||
python create_gnn_template.py --model gat --task graph_classification --output graph_classifier.py
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
TEMPLATES = {
|
||||
'node_classification': {
|
||||
'gcn': '''import torch
|
||||
import torch.nn.functional as F
|
||||
from torch_geometric.nn import GCNConv
|
||||
from torch_geometric.datasets import Planetoid
|
||||
|
||||
|
||||
class GCN(torch.nn.Module):
|
||||
"""Graph Convolutional Network for node classification."""
|
||||
|
||||
def __init__(self, num_features, hidden_channels, num_classes, num_layers=2, dropout=0.5):
|
||||
super().__init__()
|
||||
self.convs = torch.nn.ModuleList()
|
||||
|
||||
# First layer
|
||||
self.convs.append(GCNConv(num_features, hidden_channels))
|
||||
|
||||
# Hidden layers
|
||||
for _ in range(num_layers - 2):
|
||||
self.convs.append(GCNConv(hidden_channels, hidden_channels))
|
||||
|
||||
# Output layer
|
||||
self.convs.append(GCNConv(hidden_channels, num_classes))
|
||||
|
||||
self.dropout = dropout
|
||||
|
||||
def forward(self, data):
|
||||
x, edge_index = data.x, data.edge_index
|
||||
|
||||
# Apply conv layers with ReLU and dropout
|
||||
for conv in self.convs[:-1]:
|
||||
x = conv(x, edge_index)
|
||||
x = F.relu(x)
|
||||
x = F.dropout(x, p=self.dropout, training=self.training)
|
||||
|
||||
# Final layer without activation
|
||||
x = self.convs[-1](x, edge_index)
|
||||
return F.log_softmax(x, dim=1)
|
||||
|
||||
|
||||
def train(model, data, optimizer):
|
||||
"""Train the model for one epoch."""
|
||||
model.train()
|
||||
optimizer.zero_grad()
|
||||
out = model(data)
|
||||
loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
return loss.item()
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def test(model, data):
|
||||
"""Evaluate the model."""
|
||||
model.eval()
|
||||
out = model(data)
|
||||
pred = out.argmax(dim=1)
|
||||
|
||||
accs = []
|
||||
for mask in [data.train_mask, data.val_mask, data.test_mask]:
|
||||
correct = (pred[mask] == data.y[mask]).sum()
|
||||
accs.append(int(correct) / int(mask.sum()))
|
||||
|
||||
return accs
|
||||
|
||||
|
||||
def main():
|
||||
# Load dataset
|
||||
dataset = Planetoid(root='/tmp/Cora', name='Cora')
|
||||
data = dataset[0]
|
||||
|
||||
# Create model
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
model = GCN(
|
||||
num_features=dataset.num_features,
|
||||
hidden_channels=64,
|
||||
num_classes=dataset.num_classes,
|
||||
num_layers=3,
|
||||
dropout=0.5
|
||||
).to(device)
|
||||
data = data.to(device)
|
||||
|
||||
# Setup optimizer
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
|
||||
|
||||
# Training loop
|
||||
print("Training GCN model...")
|
||||
best_val_acc = 0
|
||||
for epoch in range(1, 201):
|
||||
loss = train(model, data, optimizer)
|
||||
train_acc, val_acc, test_acc = test(model, data)
|
||||
|
||||
if val_acc > best_val_acc:
|
||||
best_val_acc = val_acc
|
||||
best_test_acc = test_acc
|
||||
|
||||
if epoch % 10 == 0:
|
||||
print(f'Epoch {epoch:03d}, Loss: {loss:.4f}, '
|
||||
f'Train: {train_acc:.4f}, Val: {val_acc:.4f}, Test: {test_acc:.4f}')
|
||||
|
||||
print(f'\\nBest Test Accuracy: {best_test_acc:.4f}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
''',
|
||||
|
||||
'gat': '''import torch
|
||||
import torch.nn.functional as F
|
||||
from torch_geometric.nn import GATConv
|
||||
from torch_geometric.datasets import Planetoid
|
||||
|
||||
|
||||
class GAT(torch.nn.Module):
|
||||
"""Graph Attention Network for node classification."""
|
||||
|
||||
def __init__(self, num_features, hidden_channels, num_classes, heads=8, dropout=0.6):
|
||||
super().__init__()
|
||||
|
||||
self.conv1 = GATConv(num_features, hidden_channels, heads=heads, dropout=dropout)
|
||||
self.conv2 = GATConv(hidden_channels * heads, num_classes, heads=1,
|
||||
concat=False, dropout=dropout)
|
||||
|
||||
self.dropout = dropout
|
||||
|
||||
def forward(self, data):
|
||||
x, edge_index = data.x, data.edge_index
|
||||
|
||||
x = F.dropout(x, p=self.dropout, training=self.training)
|
||||
x = F.elu(self.conv1(x, edge_index))
|
||||
x = F.dropout(x, p=self.dropout, training=self.training)
|
||||
x = self.conv2(x, edge_index)
|
||||
|
||||
return F.log_softmax(x, dim=1)
|
||||
|
||||
|
||||
def train(model, data, optimizer):
|
||||
"""Train the model for one epoch."""
|
||||
model.train()
|
||||
optimizer.zero_grad()
|
||||
out = model(data)
|
||||
loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
return loss.item()
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def test(model, data):
|
||||
"""Evaluate the model."""
|
||||
model.eval()
|
||||
out = model(data)
|
||||
pred = out.argmax(dim=1)
|
||||
|
||||
accs = []
|
||||
for mask in [data.train_mask, data.val_mask, data.test_mask]:
|
||||
correct = (pred[mask] == data.y[mask]).sum()
|
||||
accs.append(int(correct) / int(mask.sum()))
|
||||
|
||||
return accs
|
||||
|
||||
|
||||
def main():
|
||||
# Load dataset
|
||||
dataset = Planetoid(root='/tmp/Cora', name='Cora')
|
||||
data = dataset[0]
|
||||
|
||||
# Create model
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
model = GAT(
|
||||
num_features=dataset.num_features,
|
||||
hidden_channels=8,
|
||||
num_classes=dataset.num_classes,
|
||||
heads=8,
|
||||
dropout=0.6
|
||||
).to(device)
|
||||
data = data.to(device)
|
||||
|
||||
# Setup optimizer
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)
|
||||
|
||||
# Training loop
|
||||
print("Training GAT model...")
|
||||
best_val_acc = 0
|
||||
for epoch in range(1, 201):
|
||||
loss = train(model, data, optimizer)
|
||||
train_acc, val_acc, test_acc = test(model, data)
|
||||
|
||||
if val_acc > best_val_acc:
|
||||
best_val_acc = val_acc
|
||||
best_test_acc = test_acc
|
||||
|
||||
if epoch % 10 == 0:
|
||||
print(f'Epoch {epoch:03d}, Loss: {loss:.4f}, '
|
||||
f'Train: {train_acc:.4f}, Val: {val_acc:.4f}, Test: {test_acc:.4f}')
|
||||
|
||||
print(f'\\nBest Test Accuracy: {best_test_acc:.4f}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
''',
|
||||
|
||||
'graphsage': '''import torch
|
||||
import torch.nn.functional as F
|
||||
from torch_geometric.nn import SAGEConv
|
||||
from torch_geometric.datasets import Planetoid
|
||||
|
||||
|
||||
class GraphSAGE(torch.nn.Module):
|
||||
"""GraphSAGE for node classification."""
|
||||
|
||||
def __init__(self, num_features, hidden_channels, num_classes, num_layers=2, dropout=0.5):
|
||||
super().__init__()
|
||||
self.convs = torch.nn.ModuleList()
|
||||
|
||||
self.convs.append(SAGEConv(num_features, hidden_channels))
|
||||
for _ in range(num_layers - 2):
|
||||
self.convs.append(SAGEConv(hidden_channels, hidden_channels))
|
||||
self.convs.append(SAGEConv(hidden_channels, num_classes))
|
||||
|
||||
self.dropout = dropout
|
||||
|
||||
def forward(self, data):
|
||||
x, edge_index = data.x, data.edge_index
|
||||
|
||||
for conv in self.convs[:-1]:
|
||||
x = conv(x, edge_index)
|
||||
x = F.relu(x)
|
||||
x = F.dropout(x, p=self.dropout, training=self.training)
|
||||
|
||||
x = self.convs[-1](x, edge_index)
|
||||
return F.log_softmax(x, dim=1)
|
||||
|
||||
|
||||
def train(model, data, optimizer):
|
||||
model.train()
|
||||
optimizer.zero_grad()
|
||||
out = model(data)
|
||||
loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
return loss.item()
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def test(model, data):
|
||||
model.eval()
|
||||
out = model(data)
|
||||
pred = out.argmax(dim=1)
|
||||
|
||||
accs = []
|
||||
for mask in [data.train_mask, data.val_mask, data.test_mask]:
|
||||
correct = (pred[mask] == data.y[mask]).sum()
|
||||
accs.append(int(correct) / int(mask.sum()))
|
||||
|
||||
return accs
|
||||
|
||||
|
||||
def main():
|
||||
dataset = Planetoid(root='/tmp/Cora', name='Cora')
|
||||
data = dataset[0]
|
||||
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
model = GraphSAGE(
|
||||
num_features=dataset.num_features,
|
||||
hidden_channels=64,
|
||||
num_classes=dataset.num_classes,
|
||||
num_layers=2,
|
||||
dropout=0.5
|
||||
).to(device)
|
||||
data = data.to(device)
|
||||
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
|
||||
|
||||
print("Training GraphSAGE model...")
|
||||
best_val_acc = 0
|
||||
for epoch in range(1, 201):
|
||||
loss = train(model, data, optimizer)
|
||||
train_acc, val_acc, test_acc = test(model, data)
|
||||
|
||||
if val_acc > best_val_acc:
|
||||
best_val_acc = val_acc
|
||||
best_test_acc = test_acc
|
||||
|
||||
if epoch % 10 == 0:
|
||||
print(f'Epoch {epoch:03d}, Loss: {loss:.4f}, '
|
||||
f'Train: {train_acc:.4f}, Val: {val_acc:.4f}, Test: {test_acc:.4f}')
|
||||
|
||||
print(f'\\nBest Test Accuracy: {best_test_acc:.4f}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
''',
|
||||
},
|
||||
|
||||
'graph_classification': {
|
||||
'gin': '''import torch
|
||||
import torch.nn.functional as F
|
||||
from torch_geometric.nn import GINConv, global_add_pool
|
||||
from torch_geometric.datasets import TUDataset
|
||||
from torch_geometric.loader import DataLoader
|
||||
|
||||
|
||||
class GIN(torch.nn.Module):
|
||||
"""Graph Isomorphism Network for graph classification."""
|
||||
|
||||
def __init__(self, num_features, hidden_channels, num_classes, num_layers=3, dropout=0.5):
|
||||
super().__init__()
|
||||
|
||||
self.convs = torch.nn.ModuleList()
|
||||
self.batch_norms = torch.nn.ModuleList()
|
||||
|
||||
# Create MLP for first layer
|
||||
nn = torch.nn.Sequential(
|
||||
torch.nn.Linear(num_features, hidden_channels),
|
||||
torch.nn.ReLU(),
|
||||
torch.nn.Linear(hidden_channels, hidden_channels)
|
||||
)
|
||||
self.convs.append(GINConv(nn))
|
||||
self.batch_norms.append(torch.nn.BatchNorm1d(hidden_channels))
|
||||
|
||||
# Hidden layers
|
||||
for _ in range(num_layers - 2):
|
||||
nn = torch.nn.Sequential(
|
||||
torch.nn.Linear(hidden_channels, hidden_channels),
|
||||
torch.nn.ReLU(),
|
||||
torch.nn.Linear(hidden_channels, hidden_channels)
|
||||
)
|
||||
self.convs.append(GINConv(nn))
|
||||
self.batch_norms.append(torch.nn.BatchNorm1d(hidden_channels))
|
||||
|
||||
# Output MLP
|
||||
self.lin = torch.nn.Linear(hidden_channels, num_classes)
|
||||
self.dropout = dropout
|
||||
|
||||
def forward(self, data):
|
||||
x, edge_index, batch = data.x, data.edge_index, data.batch
|
||||
|
||||
for conv, batch_norm in zip(self.convs, self.batch_norms):
|
||||
x = conv(x, edge_index)
|
||||
x = batch_norm(x)
|
||||
x = F.relu(x)
|
||||
x = F.dropout(x, p=self.dropout, training=self.training)
|
||||
|
||||
# Global pooling
|
||||
x = global_add_pool(x, batch)
|
||||
|
||||
# Output layer
|
||||
x = self.lin(x)
|
||||
return F.log_softmax(x, dim=1)
|
||||
|
||||
|
||||
def train(model, loader, optimizer, device):
|
||||
"""Train the model for one epoch."""
|
||||
model.train()
|
||||
total_loss = 0
|
||||
|
||||
for data in loader:
|
||||
data = data.to(device)
|
||||
optimizer.zero_grad()
|
||||
out = model(data)
|
||||
loss = F.nll_loss(out, data.y)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
total_loss += loss.item() * data.num_graphs
|
||||
|
||||
return total_loss / len(loader.dataset)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def test(model, loader, device):
|
||||
"""Evaluate the model."""
|
||||
model.eval()
|
||||
correct = 0
|
||||
|
||||
for data in loader:
|
||||
data = data.to(device)
|
||||
out = model(data)
|
||||
pred = out.argmax(dim=1)
|
||||
correct += (pred == data.y).sum().item()
|
||||
|
||||
return correct / len(loader.dataset)
|
||||
|
||||
|
||||
def main():
|
||||
# Load dataset
|
||||
dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES')
|
||||
print(f"Dataset: {dataset}")
|
||||
print(f"Number of graphs: {len(dataset)}")
|
||||
print(f"Number of features: {dataset.num_features}")
|
||||
print(f"Number of classes: {dataset.num_classes}")
|
||||
|
||||
# Shuffle and split
|
||||
dataset = dataset.shuffle()
|
||||
train_dataset = dataset[:int(len(dataset) * 0.8)]
|
||||
test_dataset = dataset[int(len(dataset) * 0.8):]
|
||||
|
||||
# Create data loaders
|
||||
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
|
||||
test_loader = DataLoader(test_dataset, batch_size=32)
|
||||
|
||||
# Create model
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
model = GIN(
|
||||
num_features=dataset.num_features,
|
||||
hidden_channels=64,
|
||||
num_classes=dataset.num_classes,
|
||||
num_layers=3,
|
||||
dropout=0.5
|
||||
).to(device)
|
||||
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
|
||||
|
||||
# Training loop
|
||||
print("\\nTraining GIN model...")
|
||||
for epoch in range(1, 101):
|
||||
loss = train(model, train_loader, optimizer, device)
|
||||
train_acc = test(model, train_loader, device)
|
||||
test_acc = test(model, test_loader, device)
|
||||
|
||||
if epoch % 10 == 0:
|
||||
print(f'Epoch {epoch:03d}, Loss: {loss:.4f}, '
|
||||
f'Train Acc: {train_acc:.4f}, Test Acc: {test_acc:.4f}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
''',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def generate_template(model_type: str, task: str, output_path: str):
|
||||
"""Generate a GNN template file."""
|
||||
if task not in TEMPLATES:
|
||||
raise ValueError(f"Unknown task: {task}. Available: {list(TEMPLATES.keys())}")
|
||||
|
||||
if model_type not in TEMPLATES[task]:
|
||||
raise ValueError(f"Model {model_type} not available for task {task}. "
|
||||
f"Available: {list(TEMPLATES[task].keys())}")
|
||||
|
||||
template = TEMPLATES[task][model_type]
|
||||
|
||||
# Write to file
|
||||
output_file = Path(output_path)
|
||||
output_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(output_file, 'w') as f:
|
||||
f.write(template)
|
||||
|
||||
print(f"✓ Generated {model_type.upper()} template for {task}")
|
||||
print(f" Saved to: {output_path}")
|
||||
print(f"\\nTo run the template:")
|
||||
print(f" python {output_path}")
|
||||
|
||||
|
||||
def list_templates():
|
||||
"""List all available templates."""
|
||||
print("Available GNN Templates")
|
||||
print("=" * 50)
|
||||
for task, models in TEMPLATES.items():
|
||||
print(f"\\n{task.upper()}")
|
||||
print("-" * 50)
|
||||
for model in models.keys():
|
||||
print(f" - {model}")
|
||||
print()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate GNN model templates",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
python create_gnn_template.py --model gcn --task node_classification --output gcn_model.py
|
||||
python create_gnn_template.py --model gin --task graph_classification --output gin_model.py
|
||||
python create_gnn_template.py --list
|
||||
"""
|
||||
)
|
||||
|
||||
parser.add_argument('--model', type=str,
|
||||
help='Model type (gcn, gat, graphsage, gin)')
|
||||
parser.add_argument('--task', type=str,
|
||||
help='Task type (node_classification, graph_classification)')
|
||||
parser.add_argument('--output', type=str, default='gnn_model.py',
|
||||
help='Output file path (default: gnn_model.py)')
|
||||
parser.add_argument('--list', action='store_true',
|
||||
help='List all available templates')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.list:
|
||||
list_templates()
|
||||
return
|
||||
|
||||
if not args.model or not args.task:
|
||||
parser.print_help()
|
||||
print("\\n" + "=" * 50)
|
||||
list_templates()
|
||||
return
|
||||
|
||||
try:
|
||||
generate_template(args.model, args.task, args.output)
|
||||
except ValueError as e:
|
||||
print(f"Error: {e}")
|
||||
print("\\nUse --list to see available templates")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user