THUDM/CogDL项目中的模型训练教程

THUDM/CogDL项目中的模型训练教程

【免费下载链接】CogDL CogDL: A Comprehensive Library for Graph Deep Learning (WWW 2023) 项目地址: https://gitcode.com/gh_mirrors/co/CogDL

自定义模型训练逻辑

在THUDM/CogDL项目中，开发者可以灵活地实现自定义的图神经网络(GNN)训练逻辑。该项目不仅提供了丰富的预定义模型和数据集，还支持用户根据特定需求构建自己的训练流程。

以下是一个完整的自定义GNN模型训练示例，展示了如何在Cora数据集上实现一个简单的两层GCN网络：

import torch
import torch.nn as nn
import torch.nn.functional as F
from cogdl import experiment
from cogdl.datasets import build_dataset_from_name
from cogdl.layers import GCNLayer
from cogdl.models import BaseModel

class CustomGNN(BaseModel):
    def __init__(self, in_feats, hidden_size, out_feats, dropout):
        super(CustomGNN, self).__init__()
        self.conv1 = GCNLayer(in_feats, hidden_size)
        self.conv2 = GCNLayer(hidden_size, out_feats)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, graph):
        graph.sym_norm()
        h = graph.x
        h = F.relu(self.conv1(graph, self.dropout(h)))
        h = self.conv2(graph, self.dropout(h))
        return F.log_softmax(h, dim=1)

if __name__ == "__main__":
    # 加载Cora数据集
    dataset = build_dataset_from_name("cora")[0]
    
    # 初始化模型
    model = CustomGNN(
        in_feats=dataset.num_features,
        hidden_size=64,
        out_feats=dataset.num_classes,
        dropout=0.1
    )
    
    # 设置优化器
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=5e-3)
    
    # 训练循环
    model.train()
    for epoch in range(300):
        optimizer.zero_grad()
        out = model(dataset)
        loss = F.nll_loss(out[dataset.train_mask], dataset.y[dataset.train_mask])
        loss.backward()
        optimizer.step()
    
    # 评估模型
    model.eval()
    _, pred = model(dataset).max(dim=1)
    correct = float(pred[dataset.test_mask].eq(dataset.y[dataset.test_mask]).sum().item())
    acc = correct / dataset.test_mask.sum().item()
    print(f'测试集准确率: {acc:.6f}')

这个示例展示了如何从零开始构建一个完整的GNN训练流程，包括模型定义、数据加载、训练循环和性能评估。

统一训练器架构

THUDM/CogDL项目设计了一个高度模块化的统一训练器架构，将GNN训练过程分解为四个解耦的组件：

1. 模型(Model)：定义GNN的核心计算逻辑
2. 模型包装器(Model Wrapper)：处理训练和测试步骤
3. 数据集(Dataset)：提供图数据
4. 数据包装器(Data Wrapper)：构建数据加载器

这种设计使得不同组件可以自由组合，极大提高了代码的复用性和灵活性。例如，开发者可以轻松地将原本设计用于全图训练的模型（如GCN）应用于大规模图数据，只需更换数据包装器为邻居采样或图聚类方式，而无需修改模型本身。

实验API简化流程

对于快速实验和原型开发，CogDL提供了更简洁的实验API：

from cogdl import experiment

# 最简单的方式 - 指定模型和数据集名称
experiment(model="gcn", dataset="cora")

# 更灵活的方式 - 分别构建组件
from cogdl.datasets import build_dataset
from cogdl.models import build_model
from cogdl.options import get_default_args 

args = get_default_args(model="gcn", dataset="cora")
dataset = build_dataset(args)
model = build_model(args)
experiment(model=model, dataset=dataset)

模型保存与加载

CogDL提供了便捷的模型保存和恢复功能：

# 保存训练好的模型
experiment(model="gcn", dataset="cora", checkpoint_path="gcn_cora.pt")

# 从检查点恢复训练
experiment(
    model="gcn", 
    dataset="cora", 
    checkpoint_path="gcn_cora.pt", 
    resume_training=True
)

节点嵌入保存与评估

对于图表示学习任务，CogDL支持保存节点嵌入并评估其在下游任务中的表现：

# 保存节点嵌入
experiment(
    model="prone", 
    dataset="blogcatalog", 
    save_emb_path="./embeddings/prone_blog.npy"
)

# 评估已保存的嵌入
experiment(
    model="prone",
    dataset="blogcatalog",
    load_emb_path="./embeddings/prone_blog.npy",
    num_shuffle=5,  # 重复实验次数
    training_percents=[0.1, 0.5, 0.9]  # 不同训练比例
)

通过这种模块化设计，THUDM/CogDL项目为图神经网络的研究和开发提供了高效、灵活的工具，大大降低了实验复杂度和代码重复率。

【免费下载链接】CogDL CogDL: A Comprehensive Library for Graph Deep Learning (WWW 2023) 项目地址: https://gitcode.com/gh_mirrors/co/CogDL