预训练大模型在新冠药物研发中崭露头角——首届全球AI药物研发算法大赛基线模型详解

赛题背景

首届全球AI药物研发算法大赛，是由清华大学药学院、百度飞桨、百度智能云和临港实验室共同主办，并得到中国药学会等单位大力支持的一项全球性技术创新大赛，旨在借助百度飞桨在生物计算方向上的算法优势，通过AI Studio平台，探索AI+药学领域前沿技术，挖掘和培育优秀人才，诚邀全球范围内生物计算、人工智能等相关专业的高校师生、企业、科研机构及开发者参赛。

这次大赛聚焦于新冠小分子药物等热点议题，旨在鼓励参赛者利用人工智能方法，发掘治疗新冠病毒的潜在药物。 参赛者可以利用深度学习、分子对接等方法，预测和评估小分子与主蛋白酶之间的相互作用，以及小分子在细胞中抑制病毒复制的潜力，挖掘潜在的药物候选物。通过本次大赛，我们希望能够推动药物研发领域的创新，为未来的疾病治疗和防控工作提供有力的支持。

基线模型解析

数据探索和预处理

# 导入数据集
import pandas as pd
train_df = pd.read_csv('data/data220139/train.csv')
print(f'len of train_df is {len(train_df)}')
train_df.head()

 SMILES转3D图

# 使用分子力场将smiles转化为3d分子图，并保存为smiles_to_graph_dict.pkl文件
from threading import Thread, Lock
from pahelix.utils.compound_tools import mol_to_geognn_graph_data_MMFF3d
from rdkit.Chem import AllChem

mutex = Lock()  # 互斥锁，防止多个线程同时修改某一文件或某一全局变量，引发未知错误

def calculate_3D_structure_(smiles_list):
    n = len(smiles_list)
    global p
    index = 0
    while True:
        mutex.acquire()  # 获取锁
        if p >= n:
            mutex.release()
            break
        index = p        # p指针指向的位置为当前线程要处理的smiles
        smiles = smiles_list[index]
        print(index, ':', round(index / n * 100, 2), '%', smiles)
        p += 1           # 修改全局变量p
        mutex.release()  # 释放锁
        try:
            molecule = AllChem.MolFromSmiles(smiles)
            molecule_graph = mol_to_geognn_graph_data_MMFF3d(molecule)  # 根据分子力场生成3d分子图
        except:
            print("Invalid smiles!", smiles)
            mutex.acquire()
            with open('work/invalid_smiles.txt', 'a') as f:
                # 生成失败的smiles写入txt文件保存在该目录下
                f.write(str(smiles) + '\n')
            mutex.release()
            continue

        global smiles_to_graph_dict
        mutex.acquire()   # 获取锁
        smiles_to_graph_dict[smiles] = molecule_graph
        mutex.release()   # 释放锁

for mode in ['train', 'test']:
# for mode in ['test']:
    if mode == 'train':
        smiles_list = train_df["SMILES"].tolist()
    else:
        smiles_list = test_df["SMILES"].tolist()
    global smiles_to_graph_dict
    smiles_to_graph_dict = {}
    global p              # p为全局指针，指向即将要处理的smiles
    p = 0
    thread_count = 12      # 线程数。一般根据当前运行环境下cpu的核数来设定
    threads = []
    for i in range(thread_count):
        threads.append(Thread(target=calculate_3D_structure_, args=(smiles_list, )))
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    pkl.dump(smiles_to_graph_dict, open(f'work/{mode}_smiles_to_graph_dict.pkl', 'wb'))
    print(f'{mode} is Done!')

这段代码是用于将SMILES（Simplified Molecular Input Line Entry System）表示的分子转换为3D分子图，并将结果保存为一个字典文件。SMILES是一种表示分子结构的简单文本字符串。这段代码使用了多线程来加速处理过程。下面是代码的逐行解释：

1. 导入所需的库和模块。

   • Thread 和 Lock 用于多线程处理。
   • mol_to_geognn_graph_data_MMFF3d 用于将分子转换为3D图。
   • AllChem 来自rdkit.Chem，用于处理化学信息。

2. 创建一个互斥锁，防止多个线程同时修改文件或全局变量。

3. 定义一个函数 calculate_3D_structure_，它接受一个SMILES字符串列表，并计算每个SMILES的3D分子图。

   • 使用全局变量 p 作为指针，指向要处理的SMILES。
   • 使用互斥锁来同步对全局变量的访问。
   • 对每个SMILES，尝试将其转换为3D分子图。
   • 如果转换失败，将无效的SMILES写入一个文本文件。
   • 如果转换成功，将结果添加到全局字典 smiles_to_graph_dict 中。

4. 在主程序中，对训练和测试数据集执行上述转换。

   • 初始化空字典 smiles_to_graph_dict 和指针 p。
   • 创建多个线程来并行处理SMILES列表。
   • 启动所有线程，并等待它们完成。
   • 使用pickle库将结果字典保存为文件。

这段代码的主要目的是通过多线程加速将化学分子（以SMILES形式表示）转换为3D图的过程，并将结果保存为字典文件。这可能是在化学或药物研究领域的一个预处理步骤，以便于后续的分析和建模。

模型架构

class ADMET(nn.Layer):
    def __init__(self):
        super(ADMET, self).__init__()
        compound_encoder_config = json.load(open('GEM/model_configs/geognn_l8.json', 'r'))  
        self.encoder = GeoGNNModel(compound_encoder_config) 
        self.encoder.set_state_dict(pdl.load("GEM/weight/class.pdparams")) 
        # GEM编码器输出的图特征为32维向量, 因此mlp的输入维度为32
        self.mlp = nn.Sequential(       
            nn.Linear(32, 32, weight_attr=nn.initializer.KaimingNormal()),  
            nn.ReLU(),
            nn.Linear(32, 32, weight_attr=nn.initializer.KaimingNormal()),
            nn.ReLU(),
            nn.Linear(32, 32, weight_attr=nn.initializer.KaimingNormal()),
            nn.ReLU(),
            nn.Linear(32, 2, weight_attr=nn.initializer.KaimingNormal()),
        )

    def forward(self, atom_bond_graph, bond_angle_graph):
        node_repr, edge_repr, graph_repr = self.encoder(atom_bond_graph.tensor(), bond_angle_graph.tensor())
        return self.mlp(graph_repr)

训练模型

def trial(model_version, model, batch_size, criterion, scheduler, opt):
    # 创建dataloader
    train_data_loader, valid_data_loader = get_data_loader(mode='train', batch_size=batch_size)   

    current_best_metric = -1e10
    max_bearable_epoch = 50    # 设置早停的轮数为50，若连续50轮内验证集的评价指标没有提升，则停止训练
    current_best_epoch = 0

    train_metric_list = []     # 记录训练过程中各指标的变化情况
    valid_metric_list = []

    for epoch in range(800):   # 设置最多训练800轮
        model.train()
        for (atom_bond_graph, bond_angle_graph, label_true_batch) in train_data_loader:
            label_predict_batch = model(atom_bond_graph, bond_angle_graph)
            label_true_batch = pdl.to_tensor(label_true_batch, dtype=pdl.int64, place=pdl.CUDAPlace(0))
            loss = criterion(label_predict_batch, label_true_batch)
            loss.backward()   # 反向传播
            opt.step()   # 更新参数
            opt.clear_grad()
        scheduler.step()   # 更新学习率

        # 评估模型在训练集、验证集的表现
        metric_train = evaluate(model, train_data_loader)
        metric_valid = evaluate(model, valid_data_loader)

        train_metric_list.append(metric_train)
        valid_metric_list.append(metric_valid)

        score = round((metric_valid['ap'] + metric_valid['auc']) / 2, 4)

        if score > current_best_metric:
            # 保存score最大时的模型权重
            current_best_metric = score
            current_best_epoch = epoch
            pdl.save(model.state_dict(), "weight/" + model_version + ".pkl")
        print("=========================================================")
        print("Epoch", epoch)
        pprint(("Train", metric_train))
        pprint(("Validate", metric_valid))
        print('current_best_epoch', current_best_epoch, 'current_best_metric', current_best_metric)
        if epoch > current_best_epoch + max_bearable_epoch:
            break

    return train_metric_list, valid_metric_list

最终效果

如有想组队的伙伴，请加入：