PyTorch深度学习实战 |手算GCN (图神经网络)模型

pip install rdkit

pip install torch_geometric

# 假设 PyTorch, RDKit, PyTorch Geometric (PyG) 库已安装
import torch
from rdkit import Chem
from torch_geometric.utils.smiles import from_smiles  # PyG中用于SMILES转换的实用函数 (或使用更早版本的'from_rdkit')
# --- 1. 定义和转换 ---
smiles_ethanol = "CCO" 
# 使用 PyG 的封装函数，一步完成解析、特征提取和图结构构建
# 这个函数内部自动完成了原子特征编码、键索引构建、以及 PyTorch 张量转换。
ethanol_data = from_smiles(smiles_ethanol)
# --- 2. 展示结果 ---
print("=" * 40)
print(f"乙醇分子 SMILES: {smiles_ethanol}")
print("PyG Data 对象结构 (封装结果)")
print("=" * 40)
# PyG Data 对象概览
print(ethanol_data) 
print("\n--- 关键张量尺寸分析 ---")
print(f"节点特征矩阵 (x): {ethanol_data.x.shape}")
print(f"邻接信息 (edge_index): {ethanol_data.edge_index.shape}")
print("-" * 40)

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, global_mean_pool
from torch_geometric.data import Data
import sys
# 尝试导入 RDKit 和 PyG 转换工具
from rdkit import Chem
from torch_geometric.utils.smiles import from_smiles
# --- 1. 真实数据准备 (使用 PyG 封装函数) ---
smiles_ethanol = "CCO" 
# 一步转换：生成包含所有 9 个原子（C, C, O, 6H）的图结构
ethanol_data = from_smiles(smiles_ethanol)
 # 修复：确保节点特征是浮点类型 (解决 RuntimeError)
ethanol_data.x = ethanol_data.x.float()
# 生成 Batch Tensor：9 个节点都属于同一个图 (batch size=1)
N_NODES = ethanol_data.x.shape[0]
ethanol_data.batch = torch.zeros(N_NODES, dtype=torch.long)
# --- 2. 展示输入数据结构 ---
print("=" * 60)
print(f"乙醇分子 SMILES: {smiles_ethanol}")
print("PyG Data 对象结构 (GCN 模型输入 - 真实配置)")
print("=" * 60)
print(f"节点数 N: {ethanol_data.x.shape[0]}")
print(f"节点特征矩阵 (x): {ethanol_data.x.shape}")
print(f"邻接信息 (edge_index): {ethanol_data.edge_index.shape}")
print("-" * 60)
# --- 3. 定义 GCN 模型类 (SimpleGCN) ---
class SimpleGCN(torch.nn.Module):
    def __init__(self, num_node_features, hidden_channels, num_classes):
        super().__init__()
        self.conv1 = GCNConv(num_node_features, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, hidden_channels)
        self.lin = torch.nn.Linear(hidden_channels, num_classes)
    def forward(self, x, edge_index, batch):
        print(f"\n[A] 初始输入 x (H(0)): {x.shape}")
        # GCN 层 1
        x = self.conv1(x, edge_index)
        print(f"[B] GCNConv 1 输出 (H(1)): {x.shape}")
        x = F.relu(x)
        # GCN 层 2
        x = self.conv2(x, edge_index)
        print(f"[C] GCNConv 2 输出 (H(2)): {x.shape}")
        x = F.relu(x)
        # 全局读出/池化层
        x = global_mean_pool(x, batch)
        print(f"[D] Global Mean Pool 输出: {x.shape} <--- **图级特征**")
        # 线性分类层
        x = self.lin(x)
        print(f"[E] 最终分类层输出: {x.shape} <--- **预测结果**")
        return x
# --- 4. 模型实例化与运行 ---
# 定义模型参数
INPUT_DIM = ethanol_data.x.shape[1] # 自动获取真实/模拟的特征维度 (例如：11)
HIDDEN_DIM = 16 
OUTPUT_DIM = 1 
# 实例化模型
model = SimpleGCN(
    num_node_features=INPUT_DIM, 
    hidden_channels=HIDDEN_DIM, 
    num_classes=OUTPUT_DIM
)
# 执行前向传播
print("=" * 60)
print(f"【Simple GCN 前向传播过程（隐藏层维度 D_hidden={HIDDEN_DIM}）】")
print("=" * 60)
output = model(
    ethanol_data.x, 
    ethanol_data.edge_index, 
    ethanol_data.batch
)