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

本文链接：https://blog.csdn.net/qq_40671063/article/details/130261515

这段代码展示了如何使用Python和NumPy库生成训练、验证和测试数据集的ID，并通过随机化处理。接着定义了一个函数`sample_mask`用于创建mask，最后应用该函数为不同数据集生成对应的mask。这段代码适用于机器学习模型的数据预处理阶段。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

首先生成训练集、验证集和测试集的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截图