pyg库自定义图数据集
PyG 提供了torch_geometric.data.Data用于构建图，包括 5 个属性，每一个属性都不是必须的，可以为空。
Data(x, edge_index, edge_attr, y）
x: 存储每个节点的特征，形状是[num_nodes, num_node_features]，一般是float tensor。
edge_index: 用于存储节点之间的边，形状是 [2, num_edges]，一般是long tensor。
edge_attr: 表示边属性，shape: [num_edges, num_edge_features]
y: 存储样本标签。如果是每个节点都有标签，那么形状是[num_nodes, *]；如果是整张图只有一个标签，那么形状是[1, *]，一般是long tensor。
edge_attr: 存储边的特征。形状是[num_edges, num_edge_features]
pos: 存储节点的坐标，形状是[num_nodes, num_dimensions]
那么好，我们考虑最简单的图，如果要构建一个pyg的Data，需要一个特征矩阵，一个边表，和一个标签。构建一个数据集呢？如果你的数据集就是一张图，就把这张图构建成一个Data保存成数据集，如果你的数据集有很多张图，你就可以创建一个列表，把这些构建成Data的图存到列表里，然后把列表保存成一个数据集就可以了。
PyG 提供了dataset可以自定义图数据集，其方式与Pytorch类似，需要继承数据集类。PyG中提供了两个数据集抽象类：
torch_geometric.data.InMemoryDataset：用于构建内存数据集（小数据集），继承自Dataset，一次性加载所有数据到内存。
torch_geometric.data.Dataset：用于构建大型数据集（非内存数据集），分次加载到内存；
其中的一些参数：
root：string，保存数据集的路径。
transform：将Data类型的数据作为输入，并返回转换后的图。数据对象将在每次访问之前进行转换。
pre_transform：将Data类型的数据作为输入，并返回转换后的图。数据对象将在保存到硬盘之前进行转换。
pre_filter：将Data类型的数据作为输入，并返回布尔值。指示数据对象是否应包含在最终的数据集中。
我们需要实现的方法：
raw_file_names()：返回原始数据集的文件名列表，若self.raw_dir中没有该列表中的文件，则会通过download()进行下载；
processed_file_names()：返回process()方法处理后的文件名列表，若self.processed_dir中没有确实该列表中的文件，则需要通过process()方法进行处理；
download()：下载原始数据集到self.raw_dir中，在自定义数据集中一般pass掉。
process()：写一个函数处理原始数据集成torch_geometric.data.Data的形式，并保存到processed_dir中，如果是图分类，还需要把多个图存成一个list。
就是说在pyg的dataset里面，数据集应该存储在self.raw_dir这个文件夹里面，如果这个路径下没有数据集，那他就会自动下载（这是pyg自带数据集的情况），如果是我们自建的数据集，下载肯定是没得下载的，所以我们就会把download()方法pass掉，然后dataset就会调用process方法生成数据集。所以自建数据集最关键的部分在process方法里，我们主要是实现这个。
此外，在我参考那篇文章里，他说如果是直接继承torch_geometric.data.Dataset，而不是InMemoryDataset和，需要多写两个函数len()和get():
len()：返回存储在 dataset 中的图的数目。
get()：根据idx获取数据，即单个Data图。
但是据我实测，继承InMemoryDataset最好也实现一下，都是分情况的，实现一下也不费劲，不然可能会出问题。
现在基础结束了，我们直接举例说明，在Data的基础上构建一个InMemoryDataset。我们要做的是一个图级任务，就是数据集里有很多图，构建每一张图，我们要拿到它的特征矩阵、边表、标签（回归任务·，标签是个数）。特征矩阵比较大就不展示了，边表是这样子的：
对于特征矩阵和边表，我们都要转成tensor的向量然后构建一个pyg的图Data：
# 获得特征矩阵
print('获得特征矩阵...')
node_features = torch.tensor(df_temp.drop('id',axis=1).values)
# 获得边表
print('获得边表...')
edge_list = self.get_edge_list(df_temp)
# 获得图的标签
label = float(df_train_scores[df_train_scores['id'] == id]['score'])
print('获得标签（score）：',label)
print('创建data...')
data = Data(x=node_features, edge_index=edge_list, y=label)
然后把这个Data加入到一个list里，最后把整个list保存成数据集，完整代码：
# 自建图数据集
class MyDatset(InMemoryDataset):
def __init__(self, root='./data',filepath='./data', train_path='./data/train_logs.csv',train_scores_path='./data/train_scores.csv',
name='My_processes_to_quality_data', use_edge_attr=True, transform=None,
pre_transform=None, pre_filter=None):
self.name = name
self.root = root
self.filepath = filepath
self.filenames = os.listdir(filepath)
self.train_path = train_path
self.train_scores_path = train_scores_path
self.use_edge_attr = use_edge_attr
self.pre_transform = pre_transform
self.pre_filter = pre_filter
super().__init__(root, transform, pre_transform, pre_filter)
self.data = torch.load(os.path.join(self.processed_dir, f'data.pt'))
@property
def raw_dir(self):
"""默认也是self.root/raw"""
return self.filepath
@property
def processed_dir(self):
"""默认是self.root/processed"""
return os.path.join(self.root, self.name)
@property
def raw_file_names(self):
""""原始文件的文件名，如果存在则不会触发download"""
return self.filenames
@property
def processed_file_names(self):
"""处理后的文件名，如果在 processed_dir 中找到则跳过 process"""
return ['data.pt']
def download(self):
"""这里不需要下载"""
pass
def norm(self,df_train_logs):
df_train_logs_norm = (df_train_logs[['down_time', 'action_time', 'cursor_position',
'word_count']] - df_train_logs[['down_time', 'action_time', 'cursor_position',
'word_count']].min()) / (df_train_logs[['down_time', 'action_time', 'cursor_position',
'word_count']].max() - df_train_logs[['down_time', 'action_time', 'cursor_position',
'word_count']].min())
return df_train_logs_norm
def get_edge_list(self,df_temp):
df_event_id_count = df_temp['event_id'].value_counts().to_frame().reset_index().sort_values(by='count', ascending=[False])
node_id_list = df_event_id_count["event_id"].values
print('节点数量：',node_id_list.shape[0])
node_id_list_temp1 = pd.DataFrame(np.delete(node_id_list,node_id_list.shape[0]-1))
node_id_list_temp2 = pd.DataFrame(np.delete(node_id_list,0))
edge_list = pd.concat([node_id_list_temp1, node_id_list_temp2], axis=1).transpose()
edge_list = np.transpose(np.array(edge_list)).tolist()
return torch.tensor(edge_list)
def process_data(self,df_train):
# 取要删掉的部分
df_event_count = df_train_logs['up_event'].value_counts().to_frame().reset_index().sort_values(by='count', ascending=[False])
df_activity_count = df_train_logs['activity'].value_counts().to_frame().reset_index().sort_values(by='count', ascending=[False])
df_text_change_count = df_train_logs['text_change'].value_counts().to_frame().reset_index().sort_values(by='count', ascending=[False])
vocab_event = df_event_count.loc[df_event_count['count'] < 100, "up_event"].values
vocab_activity = df_activity_count.loc[df_activity_count['count'] < 2, "activity"].values
vocab_text_change = df_text_change_count.loc[df_text_change_count['count'] < 100, "text_change"].values
# 进行多列筛选，找到在原表中的位置并删除
df_bom = df_train[df_train['activity'].isin(vocab_activity) | df_train['up_event'].isin(vocab_event) | df_train['text_change'].isin(vocab_text_change)]
df_train = df_train[~df_train['id'].isin(df_bom['id'])]
# one-hot编码
df_train = df_train[['id', 'event_id', 'down_time', 'action_time', 'activity','down_event', 'text_change', 'cursor_position','word_count']]
df_train = pd.get_dummies(df_train,columns=['activity','down_event','text_change'],dtype=int)
# 归一化
df_train_logs_temp = self.norm(df_train)
df_train = df_train.drop(['down_time', 'action_time', 'cursor_position','word_count'],axis=1)
df_train = df_train.join(df_train_logs_temp)
print('处理后数据集维度：',df_train.shape)
return df_train
def process(self):
"""主程序，对原始数据进行处理"""
graphs = []
# 读出来数据和标签
df_train_scores = pd.read_csv(self.train_scores_path)
df_train_logs = pd.read_csv(self.train_path)
# 处理好数据，归一化、onehot编码等
df_train_logs = self.process_data(df_train_logs)
# 拿出来id列表
df_id_count = df_train_logs['id'].value_counts().to_frame().reset_index().sort_values(by='count', ascending=[False])
vocab_id = df_id_count["id"].values
for id in vocab_id:
print('id:',id)
df_temp = df_train_logs[df_train_logs['id'] == id]
# 获得特征矩阵
print('获得特征矩阵...')
node_features = torch.tensor(df_temp.drop('id',axis=1).values)
# 获得边表
print('获得边表...')
edge_list = self.get_edge_list(df_temp)
# 获得图的标签
label = float(df_train_scores[df_train_scores['id'] == id]['score'])
print('获得标签（score）：',label)
print('创建data...')
data = Data(x=node_features, edge_index=edge_list, y=label)
graphs.append(data)
torch.save(graphs, os.path.join(self.processed_dir, f'data.pt'))
def get(self, idx):
# data = torch.load(os.path.join(self.processed_dir, f'data.pt'))
return self.data[idx]
def len(self):
return len(self.data)
当我们第一次实例化这个dataset，会自动检测那个文件夹是否有数据集，如果没有就会执行process方法构建数据集，下次实例化就不会进行这一步了：
dataset = MyDatset()
可以直接用pyg的DataLoader读出数据：
from torch_geometric.data import DataLoader
from torch_scatter import scatter_mean
dataset = MyDatset()
# shuffle:是否打乱顺序
# train_loader = DataLoader(dataset, batch_size=8, shuffle=True,drop_last=True)
train_loader = DataLoader(dataset,batch_size=2)
for data in train_loader:
print(data)
一个简易的用于回归任务的GCN：
from torch.nn import Linear
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
from torch_geometric.nn import global_mean_pool
class GCN(torch.nn.Module):
def __init__(self, in_channels, hidden_channels_1,hidden_channels_2, hidden_channels_3,out_channels):
super(GCN, self).__init__()
torch.manual_seed(12345)
self.conv1 = GCNConv(in_channels, hidden_channels_1)
self.conv2 = GCNConv(hidden_channels_1, hidden_channels_2)
self.conv3 = GCNConv(hidden_channels_2, hidden_channels_3)
self.lin = Linear(hidden_channels_3, out_channels)
def forward(self, x, edge_index, batch):
# 1. 获得节点嵌入
x=x.to(torch.float32)
print(x.shape)
print(edge_index.shape)
x = self.conv1(x, edge_index)
x = x.relu()
x = self.conv2(x, edge_index)
x = x.relu()
x = self.conv3(x, edge_index)
# 2. 对全局的点嵌入池化，得到图嵌入
x = global_mean_pool(x, batch) # [batch_size, hidden_channels]
# 3. 回归
x = F.dropout(x, p=0.5, training=self.training)
x = self.lin(x)
return x
实例化模型和训练：
def my_RMSE(labels, res):
sum1 = 0
for i in range(len(labels)):
sum1 = sum1 + (labels.iloc[i] - res.iloc[i]) ** 2
RMSE = (sum1 / len(labels)) ** 0.5
return RMSE
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
input_dim = dataset.num_node_features
output_dim = 1
print('imput_dim:',input_dim)
model = GCN(input_dim, 64,32,16, output_dim).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
criterion = torch.nn.CrossEntropyLoss()
def train():
model.train()
for data in train_loader:
# print('进入训练')
optimizer.zero_grad()
out = model(data.x, data.edge_index, data.batch)
loss = criterion(out, data.y)
loss.backward()
optimizer.step()
def test(loader):
model.eval()
RMSE = 0
outs = []
y = []
for data in loader: # 批遍历测试集
out = model(data.x, data.edge_index, data.batch) # 一次前向传播
# print('out:',out)
# print('y:',data.y)
outs.append(out)
y.append(data.y)
return my_RMSE(y, outs)
for epoch in range(1, 120):
print('epoch:',epoch)
train()
if epoch % 10 == 0:
RMSE = test(train_loader)
print(f'Epoch: {epoch:03d}, Train RMSE: {RMSE:.4f}')
本文有关pyg库的介绍参考了：PyG自定义数据集学习笔记