CGCNN图神经网络入门实战：从分子结构预测到模型部署全流程解析

1次阅读

共计 3336 个字符，预计需要花费 9 分钟才能阅读完成。

材料科学领域中，密度泛函理论（DFT）计算虽然精度高，但计算成本巨大，单个晶体结构的计算可能需要数小时甚至数天。传统机器学习方法（如随机森林、支持向量机）虽然计算速度快，但难以有效捕捉晶体结构的复杂拓扑特征和原子间的多体相互作用。

模型	适用场景	特征处理方式	计算效率
CGCNN	晶体材料	晶体图卷积	高
SchNet	分子体系	连续滤波卷积	中
MPNN	通用分子图	消息传递机制	低

晶体图由原子节点和边组成，节点特征包括原子类型、价电子数等，边特征通常包含原子间距和化学键信息。

import numpy as np
from pymatgen import Structure

# 从 cif 文件构建晶体图
def build_crystal_graph(cif_path):
    structure = Structure.from_file(cif_path)
    # 获取邻接原子对
    neighbors = structure.get_all_neighbors(r=5.0)

    node_features = []
    edge_features = []
    edge_indices = []

    for i, site in enumerate(structure):
        # 节点特征：原子序数、价电子数等
        node_features.append([site.specie.Z, site.specie.electrons])

        # 边特征：距离、键角等
        for neighbor in neighbors[i]:
            j = neighbor.index
            distance = neighbor.distance
            edge_indices.append([i, j])
            edge_features.append([distance])

    return np.array(node_features), np.array(edge_features), np.array(edge_indices).T

CGCNN 的核心是晶体图卷积层，它同时考虑中心原子和邻居原子的相互作用：

import torch
import torch.nn as nn
import torch.nn.functional as F

class CGCNNLayer(nn.Module):
    def __init__(self, node_dim, edge_dim, out_dim):
        super().__init__()
        # 节点特征变换
        self.node_fc = nn.Linear(node_dim, out_dim)

        # 边特征变换
        self.edge_fc = nn.Sequential(nn.Linear(edge_dim, out_dim),
            nn.Sigmoid())

        # 多体交互项
        self.mlp = nn.Sequential(nn.Linear(out_dim * 2, out_dim),
            nn.Softplus())

    def forward(self, x, edge_attr, edge_index):
        row, col = edge_index

        # 节点特征传播
        out = self.node_fc(x)

        # 边特征处理
        edge_embed = self.edge_fc(edge_attr)

        # 邻居聚合
        agg = torch.zeros_like(out)
        agg = agg.index_add_(0, row, out[col] * edge_embed)

        # 结合中心节点和邻居信息
        combined = torch.cat([out, agg], dim=-1)
        return self.mlp(combined)

特征标准化：

from sklearn.preprocessing import StandardScaler

# 节点特征标准化
node_scaler = StandardScaler()
node_features = node_scaler.fit_transform(node_features)

# 边特征标准化
edge_scaler = StandardScaler()
edge_features = edge_scaler.fit_transform(edge_features)

数据集划分（80% 训练，10% 验证，10% 测试）：

from sklearn.model_selection import train_test_split

# 初始训练测试分割
train_idx, test_idx = train_test_split(range(len(data)), test_size=0.2, random_state=42)

# 验证集分割
val_idx, test_idx = train_test_split(test_idx, test_size=0.5, random_state=42)

import torch.optim as optim
from torch_geometric.data import Data, DataLoader

# 1. 数据准备
train_dataset = [...]  # 构建训练集 Data 对象
val_dataset = [...]    # 构建验证集 Data 对象

# 2. 模型定义
model = CGCNN(node_dim=node_features.shape[1], 
              edge_dim=edge_features.shape[1],
              out_dim=64)

# 3. 训练配置
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64)

# 4. 训练循环
for epoch in range(100):
    model.train()
    for batch in train_loader:
        optimizer.zero_grad()
        out = model(batch.x, batch.edge_attr, batch.edge_index)
        loss = criterion(out, batch.y)
        loss.backward()
        optimizer.step()

    # 验证
    model.eval()
    val_loss = 0
    with torch.no_grad():
        for batch in val_loader:
            out = model(batch.x, batch.edge_attr, batch.edge_index)
            val_loss += criterion(out, batch.y).item()

    print(f'Epoch {epoch}: Train Loss {loss.item():.4f}, Val Loss {val_loss/len(val_loader):.4f}')

Batch Size 选择：
小 batch（32-64）适合中等规模 GPU（如 RTX 2080）
大 batch（128-256）需要高端 GPU（如 V100）并减少更新频率

GPU 并行：

# 多 GPU 训练
model = nn.DataParallel(model)

周期性边界条件：
使用 pymatgen 正确识别晶胞
确保邻接原子包含周期性镜像
小样本过拟合：
添加 Dropout 层（p=0.1-0.3）
使用早停（patience=10-20）

ONNX 部署：

# 转换模型
torch.onnx.export(model, 
                  (x_sample, edge_attr_sample, edge_index_sample),
                  "cgcnn.onnx",
                  opset_version=11)