使用CSV文件建立PYG数据集，进行分子预测

最新推荐文章于 2024-03-07 20:32:58 发布

battle不停息

最新推荐文章于 2024-03-07 20:32:58 发布

阅读量1.4k

点赞数 5

CC 4.0 BY-SA版权

分类专栏：图深度学习实战文章标签： pytorch python 深度学习

本文链接：https://blog.csdn.net/m0_52647764/article/details/125139334

本文介绍了如何使用CSV文件创建适用于PyG的图神经网络数据集，涉及分子预测任务。首先，通过rdkit库从SMILES字符串转换为分子结构，然后提取原子和键的特征作为节点和边的属性，最终保存为PyG的Dataset格式。代码示例展示了数据处理和预处理的详细步骤。

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

使用CSV文件建立PYG数据集，进行分子预测

使用自己的分子数据集（csv，sdf，mol2）来创建适用于PyG的任务。
需要的第三方库
rdkit， pytorch， pyg，

前言

在使用图神经网络进行任务是，我们有些仅仅使用轮子，所以要使用自己的数据集来预测某些指标。所以我们就需要用自己的数据集建立适合PyG的Dataset。

一.步骤

PyG有两种方法建立数据集，一种是直接读到内存中去的InMemoryDataset，限制是你的内存大小，还有一种是建立比较大的数据集的Dataset。这次我主要向分享比较全能使用的Dataset。

代码如下（示例）：

import pandas as pd
import torch
from torch_geometric.data import Dataset, Data
import numpy as np
from scipy.sparse import coo_matrix
import os
from rdkit import Chem
from rdkit.Chem.rdmolops import GetAdjacencyMatrix


class mydataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(mydataset, self).__init__(root, transform, pre_transform)

    @property
    def raw_file_names(self):
        return 'fda1.csv'

    @property
    def processed_file_names(self):
        self.data = pd.read_csv(self.raw_paths[0]).reset_index()
        return [f'data_{i}.pt' for i in list(self.data.index)]

    def download(self):
        pass

    def process(self):
        data = pd.read_csv(self.raw_paths[0])
        smiles = data['smiles'].values.tolist()
        label = self.get_label(data['values']).tolist()[0]
        for index, (smi, y) in enumerate(zip(smiles, label)):
            mol = Chem.MolFromSmiles(smi)
            n_nodes = mol.GetNumAtoms()
            n_edges = 2 * mol.GetNumBonds()
            unrelated_smiles = 'O=O'
            unrelated_mol = Chem.MolFromSmiles(unrelated_smiles)
            n_node_features = len(self.get_atom_features(unrelated_mol.GetAtomWithIdx(0)))
            n_edge_features = len(self.get_bond_feature(unrelated_mol.GetBondBetweenAtoms(0, 1)))
            x = np.zeros((n_nodes, n_node_features))
            y = np.array(y)
            y = torch.tensor(np.array