torchdrug

# TorchDrug ## 概述 TorchDrug 是一个全面的、基于 PyTorch 的药物发现和分子科学机器学习工具箱。它将图神经网络、预训练模型和任务定义应用于分子、蛋白质和生物知识图谱，包括分子性质预测、蛋白质建模、知识图谱推理、分子生成、逆合成规划等，并内置了 40 多个精选数据集和 20 多种模型架构。 ## 何时使用此技能 当处理以下内容时，应使用此技能： **数据类型：** - SMILES 字符串或分子结构 - 蛋白质序列或 3D 结构 (PDB 文件) - 化学反应和逆合成 - 生物医学知识图谱 - 药物发现数据集 **任务：** - 预测分子性质（溶解度、毒性、活性） - 蛋白质功能或结构预测 - 药物-靶点结合预测 - 生成新的分子结构 - 规划化学合成路线 - 生物医学知识库中的链接预测 - 在科学数据上训练图神经网络 **库与集成：** - TorchDrug 是主要库 - 常与 RDKit 结合用于化学信息学 - 兼容 PyTorch 和 PyTorch Lightning - 集成了 AlphaFold 和 ESM 用于蛋白质分析 ## 入门指南 ### 安装 bash uv pip install torchdrug # 或安装可选依赖 uv pip install torchdrug[full] ### 快速示例 python from torchdrug import datasets, models, tasks from torch.utils.data import DataLoader # 加载分子数据集 dataset = datasets.BBBP("~/molecule-datasets/") train_set, valid_set, test_set = dataset.split() # 定义 GNN 模型 model = models.GIN( input_dim=dataset.node_feature_dim, hidden_dims=[256, 256, 256], edge_input_dim=dataset.edge_feature_dim, batch_norm=True, readout="mean" ) # 创建性质预测任务 task = tasks.PropertyPrediction( model, task=dataset.tasks, criterion="bce", metric=["auroc", "auprc"] ) # 使用 PyTorch 训练 optimizer = torch.optim.Adam(task.parameters(), lr=1e-3) train_loader = DataLoader(train_set, batch_size=32, shuffle=True) for epoch in range(100): for batch in train_loader: loss = task(batch) optimizer.zero_grad() loss.backward() optimizer.step() ## 核心能力 ### 1. 分子性质预测 根据结构预测分子的化学、物理和生物学性质。 **用例：** - 药物相似性和 ADMET 属性预测