赞
踩
pytorch geometric是基于pytorch框架封装的图神经网络库,相比于DGL库更加便利,符合pytorch开发规范,拥有pytorch基础,更易于使用。pytorch基础请移步pytorch训练CNN
库名 | 作用 |
---|---|
torch_geometric.nn | 封装了常用的GNN相关网络,直接可用 |
torch_geometric.data | GNN相关的数据结构,数据批量加载工具等 |
torch_geometric.datasets | 公开直接可用的数据集 |
torch_geometric.transforms | 数据转换工具类 |
torch_geometric.utils | 其他工具类 |
torch_geometric.io | 读取写入类 |
PyTorch Geometric中的一个图由的实例描述torch_geometric.data.Data,默认情况下具有以下属性:
from torch_geometric.datasets import Planetoid
dataset = Planetoid(root='/tmp/Cora', name='Cora')
>>> Cora()
import torch import torch.nn.functional as F from torch_geometric.nn import GCNConv class Net(torch.nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = GCNConv(dataset.num_node_features, 16) self.conv2 = GCNConv(16, dataset.num_classes) def forward(self, data): x, edge_index = data.x, data.edge_index x = self.conv1(x, edge_index) x = F.relu(x) x = F.dropout(x, training=self.training) x = self.conv2(x, edge_index) return F.log_softmax(x, dim=1)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().to(device)
data = dataset[0].to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
for epoch in range(200):
optimizer.zero_grad()
out = model(data)
loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
loss.backward()
optimizer.step()
_, pred = model(data).max(dim=1)
correct = int(pred[data.test_mask].eq(data.y[data.test_mask]).sum().item())
acc = correct / int(data.test_mask.sum())
print('Accuracy: {:.4f}'.format(acc))
>>> Accuracy: 0.8150
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。