生成pyg构建数据的train_mask, val_mask和test_mask

首先生成训练集、验证集和测试集的id

代码：

sample_number = len(y) #y就是label 的长度
# 这些id也可以通过随机的方式生成 
train_idx = range(int(0.7*sample_number))
val_idx = range(int(0.7*sample_number), int(0.9*sample_number))
test_idx = range(int(0.9*sample_number), int(sample_number))

# 随机方式生成
shuffled_idx = shuffle(np.array(range(len(y)), random_state=seed) # 已经被随机打乱
train_idx = shuffled_idx[:int(0.7* y.shape[0])].tolist()
val_idx = shuffled_idx[int(0.7*y.shape[0]): int(0.9*label.shape[0])].tolist()
test_idx = shuffled_idx[int(0.9*y.shape[0]):].tolist()

进行mask

mask函数：

import torch


def sample_mask(idx, l):
    """Create mask."""
    mask = torch.zeros(l)
    mask[idx] = 1
    return torch.as_tensor(mask, dtype=torch.bool)

调用上述函数

train_mask = sample_mask(train_idx, sample_number)
val_mask = sample_mask(val_idx, sample_number)
test_mask = sample_mask(test_idx, sample_number)

完整代码

大家可以试着跑一下这个代码

import torch
from sklearn.utils import shuffle
import numpy as np

def sample_mask(idx, l):
    """Create mask."""
    mask = torch.zeros(l)
    mask[idx] = 1
    return torch.as_tensor(mask, dtype=torch.bool)

y = torch.arange(100)
sample_number = len(y) #y就是label 的长度
# # 这些id也可以通过随机的方式生成 
# train_idx = range(int(0.7*sample_number))
# val_idx = range(int(0.7*sample_number), int(0.9*sample_number))
# test_idx = range(int(0.9*sample_number), int(sample_number))

seed = 10
shuffled_idx = shuffle(np.array(range(len(y))), random_state=seed) # 已经被随机打乱
train_idx = shuffled_idx[:int(0.7* y.shape[0])].tolist()
val_idx = shuffled_idx[int(0.7*y.shape[0]): int(0.9*y.shape[0])].tolist()
test_idx = shuffled_idx[int(0.9*y.shape[0]):].tolist()
train_mask = sample_mask(train_idx, sample_number)
val_mask = sample_mask(val_idx, sample_number)
test_mask = sample_mask(test_idx, sample_number)

train_mask截图