在使用PyTorch Geometric加载Cora数据集时,训练、验证和测试集的大小可以通过掩码(mask)来确定。每个掩码是一个布尔张量,用于指示哪些节点属于训练集、验证集和测试集。通过计算这些掩码中为True的元素数量,我们可以确定每个数据集的大小。
示例代码
以下是如何计算训练、验证和测试集大小的示例代码:
import torch
from torch_geometric.datasets import Planetoid
from torch_geometric.transforms import NormalizeFeatures
# 加载Cora数据集
dataset = Planetoid(root='/tmp/Cora', name='Cora', transform=NormalizeFeatures())
data = dataset[0]
print(f"data.train_mask{data.train_mask}")
print(f"data.val_mask{data.val_mask}")
print(f"data.test_mask{data.test_mask}")
# 计算训练、验证和测试集的大小
num_train = data.train_mask.sum().item()
num_val = data.val_mask.sum().item()
num_test = data.test_mask.sum().item()
print(f'Number of training nodes: {num_train}')
print(f'Number of validation nodes: {num_val}')
print(f'Number of test nodes: {num_test}')
输出解释
data.train_mask
:布尔张量,指示哪些节点用于训练。data.val_mask
:布尔张量,指示哪些节点用于验证。data.test_mask
:布尔张量,指示哪些节点用于测试。sum()
:计算布尔张量中为True的元素数量。item()
:将结果转换为Python标量。
通过这种方式,你可以轻松地计算出每个数据集的大小。
运行示例
假设你运行上述代码,输出可能类似于:
data.train_masktensor([ True, True, True, ..., False, False, False])
data.val_masktensor([False, False, False, ..., False, False, False])
data.test_masktensor([False, False, False, ..., True, True, True])
Number of training nodes: 140
Number of validation nodes: 500
Number of test nodes: 1000
这些数字是根据Cora数据集的默认分割方式得出的,表示该数据集使用140个节点进行训练,500个节点进行验证,以及1000个节点进行测试。