山东大学项目实训——个人纪录（10）

目录

项目团队DrugLLM开发团队

日期2024/6/9

本周进度

---

项目团队
DrugLLM开发团队

日期
2024/6/9

本周进度

对模型要使用的数据进行清洗。将存储在pickle文件中的图数据转换格式，并重新保存。

import json
import pickle
import torch
from torch_geometric.data import Data, Batch


def main():

    with open("../data/chebi_0.pkl", "rb") as f:
        data = pickle.load(f)
    for rec in data:
        g = rec["graph"]
        graph = Data(x=torch.asarray(g['node_feat']), edge_index=torch.asarray(g['edge_index']), edge_attr=torch.asarray(g['edge_feat']))
        rec["graph"] = graph

    with open("../data/chebi_graph.pkl", "wb") as f:
        pickle.dump(data, f)


if __name__ == "__main__":
    main()

1.载入外部数据: 首先，使用open函数以二进制读模式("rb")打开文件../data/chebi_0.pkl。

2.数据格式转换: 接着，代码遍历data列表中的每一个记录rec。对于每个记录，提取出其"graph"部分，这包含了节点特征('node_feat')、边索引('edge_index')和边特征('edge_feat')。然后，利用PyTorch库，将这些特征转换为张量格式，以便于后续的图神经网络或其他基于PyTorch的操作。

3.保存转换后的数据: 最后，使用pickle.dump(data, f)将整个转换过格式的data列表，以二进制写模式("wb")保存到新的pickle文件../data/chebi_graph.pkl中。