TorchDrug教程

TorchDrug

基于张量的图

图表

图是 TorchDrug 中使用的基本数据结构。它们可以被视为张量框架中图形处理库（例如NetworkX）的泛化，同时享受 GPU 加速和自动微分。

import torchdrug as td
from torchdrug import data

edge_list = [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 0]]
graph = data.Graph(edge_list, num_node=6)
graph.visualize()

分子

药物发现中图的一个特定实例是一个分子，它可能具有额外的化学约束和特征。例如，我们可以从SMILES字符串构造一个分子。SMILES 是一种使用线性字符串描述复杂分子结构的语言。

mol = data.Molecule.from_smiles("C1=CC=CC=C1")
mol.visualize()
print(mol.node_feature.shape)
print(mol.edge_feature.shape)

图操作

一组任意图也可以批量处理，以最大限度地利用硬件。它们也可以使用 cpu()和cuda()方法在 CPU 和 GPU 之间传输。

smiles_list = ["CCSCCSP(=S)(OC)OC", "CCOC(=O)N",
               "N(Nc1ccccc1)c2ccccc2", "NC(=O)c1cccnc1"]
mols = data.PackedMolecule.from_smiles(smiles_list)
mols.visualize()
mols = mols.cuda()
print(mols)

就像原始的 PyTorch 张量一样，图支持广泛的索引操作。典型用法包括应用节点屏蔽、边缘屏蔽或图形屏蔽。以下示例显示如何选择包含至少一个碳的边。更多图操作请看注释。

node_in, node_out, bond_type = mols.edge_list.t()
edge_mask = (mols.atom_type[node_in] == td.CARBON) | \
            (mols.atom_type[node_out] == td.CARBON)
mols = mols.edge_mask(edge_mask)
mols.visualize()