Jittor Geometric 项目
Jittor Geometric 是一个基于Jittor深度学习框架构建的图机器学习库，专门用于处理图结构数据和几何深度学习任务。
项目地址项目概述
Jittor Geometric 是一个先进的图机器学习库，建立在Jittor深度学习框架之上。它提供了丰富的图神经网络模型和几何深度学习算法，支持高效的图数据处理和模型训练。
主要特性
- 高性能图神经网络实现
- 与Jittor深度学习框架无缝集成
- 支持多种图数据结构
- 丰富的几何深度学习算法
快速入门
以下是一个使用Jittor Geometric训练GCN模型的简单示例：
# Dataset Selection
import os.path as osp
from jittor_geometric.datasets import Planetoid
import jittor_geometric.transforms as T
dataset = 'Cora'
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', dataset)
dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures())
data = dataset[0]# Data Preprocess
from jittor_geometric.ops import cootocsr,cootocsc
from jittor_geometric.nn.conv.gcn_conv import gcn_norm
edge_index, edge_weight = data.edge_index, data.edge_attr
edge_index, edge_weight = gcn_norm(
edge_index, edge_weight,v_num,
improved=False, add_self_loops=True)
with jt.no_grad():
data.csc = cootocsc(edge_index, edge_weight, v_num)
data.csr = cootocsr(edge_index, edge_weight, v_num)# Model Definition
from jittor import nn
from jittor_geometric.nn import GCNConv
class GCN(nn.Module):
def __init__(self, dataset, dropout=0.8):
super(Net, self).__init__()
self.conv1 = GCNConv(in_channels=dataset.num_features, out_channels=256,spmm=args.spmm)
self.conv2 = GCNConv(in_channels=256, out_channels=dataset.num_classes,spmm=args.spmm)
self.dropout = dropout
def execute(self):
x, csc, csr = data.x, data.csc, data.csr
x = nn.relu(self.conv1(x, csc, csr))
x = nn.dropout(x, self.dropout, is_train=self.training)
x = self.conv2(x, csc, csr)
return nn.log_softmax(x, dim=1)# Training
model = GCN(dataset)
optimizer = nn.Adam(params=model.parameters(), lr=0.001, weight_decay=5e-4)
for epoch in range(200):
model.train()
pred = model()[data.train_mask]
label = data.y[data.train_mask]
loss = nn.nll_loss(pred, label)
optimizer.step(loss)