CGCNN图神经网络实战：解决分子属性预测中的特征提取难题

1次阅读

没有评论

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

分子和晶体属性预测是材料科学中的重要任务，传统方法主要依赖手工设计的特征：

领域知识依赖性强 ：需要专家设计描述符（如价电子数、电负性等），难以覆盖复杂相互作用
泛化性差 ：固定特征工程无法适应不同预测任务，每次需重新设计
长程相互作用缺失 ：传统方法难以捕捉超过 3 个原子间距的关联

对比实验显示，在 Materials Project 数据集上：

随机森林（Random Forest）的 MAE（平均绝对误差）为 0.48 eV
支持向量机（SVM）的 MAE 为 0.52 eV
简单全连接神经网络的 MAE 为 0.41 eV

将晶体结构转化为图数据：

节点特征 ：原子类型（one-hot 编码）+ 基本属性（原子序数、质量等）

边特征 ：使用径向基函数（RBF, Radial Basis Function）处理原子间距：

# RBF 展开公式
rbf = exp(-γ*(distance - center)**2)  # γ 和 center 为可学习参数

邻接矩阵 ：设置截断半径（通常 5 -8Å），超出距离的原子不建边

采用消息传递（Message Passing）机制，每层执行：

消息生成 ：结合节点 j 的特征和边 ij 的特征
消息聚合 ：对邻居节点消息求和 / 平均
节点更新 ：用聚合结果更新节点 i 特征

数学表达式：

h_i^(l+1) = σ(W_self h_i^l + Σ_{j∈N(i)} W_edge f(r_ij) ⊙ h_j^l)

其中 f(r_ij) 是边特征变换函数，⊙表示逐元素乘法

将节点特征转化为图级表示：

加权平均池化 ：给每个原子赋予可学习权重
求和池化 ：直接求和所有节点特征

from torch_geometric.data import Data
import numpy as np

class CrystalDataset(InMemoryDataset):
    def __init__(self, root, transform=None):
        super().__init__(root, transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

    def process(self):
        # 读取 CIF 文件并构建图结构
        data_list = []
        for cif_file in glob.glob('*.cif'):
            structure = Structure.from_file(cif_file)

            # 构建节点特征（原子类型）atom_features = [get_atom_feature(atom) for atom in structure]
            x = torch.tensor(atom_features, dtype=torch.float)

            # 构建边索引和边属性
            edge_index, edge_attr = build_adjacency_matrix(structure, cutoff=5.0)

            data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
            data_list.append(data)

        torch.save(self.collate(data_list), self.processed_paths[0])

import torch.nn as nn
from torch_geometric.nn import MessagePassing

class CGCNNLayer(MessagePassing):
    def __init__(self, node_dim, edge_dim):
        super().__init__(aggr='add')
        self.message_net = nn.Sequential(nn.Linear(2*node_dim + edge_dim, node_dim),
            nn.SiLU())
        self.update_net = nn.Sequential(nn.Linear(node_dim, node_dim),
            nn.SiLU())

    def forward(self, x, edge_index, edge_attr):
        return self.propagate(edge_index, x=x, edge_attr=edge_attr)

    def message(self, x_i, x_j, edge_attr):
        # x_i: 目标节点特征 [E, node_dim]
        # x_j: 源节点特征 [E, node_dim]
        # edge_attr: 边特征 [E, edge_dim]
        return self.message_net(torch.cat([x_i, x_j, edge_attr], dim=-1))

class CGCNN(nn.Module):
    def __init__(self, num_layers=3):
        super().__init__()
        self.conv_layers = nn.ModuleList([CGCNNLayer(node_dim=64, edge_dim=32) 
            for _ in range(num_layers)
        ])
        self.pool = nn.Linear(64, 1)  # 加权全局池化

    def forward(self, data):
        x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr

        # 多层图卷积
        for conv in self.conv_layers:
            x = conv(x, edge_index, edge_attr)

        # 全局池化
        graph_embedding = self.pool(x).mean(dim=0)
        return graph_embedding

使用 SMOTE 过采样时需注意：

先在特征空间做 SMOTE 生成虚拟样本
再根据虚拟原子的特征找到最近邻的真实原子组合
用这些原子构建新的晶体图

关键参数影响：

卷积层数 ：3- 5 层效果最佳，超过 7 层会出现过平滑（over-smoothing）
截断半径 ：金属材料建议 5 -6Å，有机材料建议 4 -5Å
RBF 中心点数 ：通常 8 -12 个，覆盖 0 -5Å范围

使用 Grad-CAM 可视化重要原子位点：

model.eval()
output = model(data)
output.backward()

# 获取节点梯度
node_grad = data.x.grad  # [num_nodes, feature_dim]
importance = torch.norm(node_grad, dim=1)  # 计算重要性得分

在 Materials Project 数据集（10,000 个晶体）上的结果：