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自定义数据集学习笔记