提供了编程的基础技术教程

zl程序教程

您现在的位置是：首页 >  IT要闻

当前栏目

PyG学习笔记1-INTRODUCTION BY EXAMPLE（二）

笔记 学习
2023-03-14 22:58:04 时间

自定义 Dataset

尽管 PyG 已经包含许多有用的数据集，我们也可以通过继承torch_geometric.data.Dataset使用自己的数据集。提供 2 种不同的Dataset：


InMemoryDataset：使用这个Dataset会一次性把数据全部加载到内存中。

Dataset: 使用这个Dataset每次加载一个数据到内存中，比较常用。

我们需要在自定义的Dataset的初始化方法中传入数据存放的路径，然后 PyG 会在这个路径下再划分 2 个文件夹：


raw_dir: 存放原始数据的路径，一般是 csv、mat 等格式

processed_dir: 存放处理后的数据，一般是 pt 格式 ( 由我们重写process()方法实现)。


Transforms

transforms在计算机视觉领域是一种很常见的数据增强。PyG 有自己的transforms，输出是Data类型，输出也是Data类型。可以使用torch_geometric.transforms.Compose封装一系列的transforms。我们以 ShapeNet 数据集 (包含 17000 个 point clouds，每个 point 分类为 16 个类别的其中一个) 为例，我们可以使用transforms从 point clouds 生成最近邻图：

import torch_geometric.transforms as T
from torch_geometric.datasets import ShapeNet
dataset = ShapeNet(root='/tmp/ShapeNet', categories=['Airplane'],
                    pre_transform=T.KNNGraph(k=6))
# dataset[0]: Data(edge_index=[2, 15108], pos=[2518, 3], y=[2518])


还可以通过transform在一定范围内随机平移每个点，增加坐标上的扰动，做数据增强：

import torch_geometric.transforms as T
from torch_geometric.datasets import ShapeNet
dataset = ShapeNet(root='/tmp/ShapeNet', categories=['Airplane'],
                    pre_transform=T.KNNGraph(k=6),
                    transform=T.RandomTranslate(0.01))
# dataset[0]: Data(edge_index=[2, 15108], pos=[2518, 3], y=[2518])


模型训练

这里只是展示一个简单的 GCN 模型构造和训练过程，没有用到Dataset和DataLoader。


我们将使用一个简单的 GCN 层，并在 Cora 数据集上实验。有关 GCN 的更多内容，请查看**这篇博客**。


我们首先加载数据集：

from torch_geometric.datasets import Planetoid
dataset = Planetoid(root='/tmp/Cora', name='Cora')


然后定义 2 层的 GCN：

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)


然后训练 200 个 epochs：

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)
model.train()
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()


最后在测试集上验证了模型的准确率：

model.eval()
_, pred = model(data).max(dim=1)
correct = float (pred[data.test_mask].eq(data.y[data.test_mask]).sum().item())
acc = correct / data.test_mask.sum().item()
print('Accuracy: {:.4f}'.format(acc))


参考链接

PyG Documentation — pytorch_geometric 2.0.2 documentation (pytorch-geometric.readthedocs.io)

相关文章