Datawhale AI 夏令营---第三期（催化反应产率预测）

本文链接：https://blog.csdn.net/weixin_52520094/article/details/140769445

学习笔记

Task2：RNN建模SMILES进行反应产率预测

说明：在task1的学习中，使用RDkit工具将SMILES表示的分子转化为了Morgan Fingerprint（向量），然后使用随机森林对数据进行建模得到了我们的第一个成绩，并且通过调参改进的尝试，分数也在原先的基础上得到较大的提升。接下来本次的学习任务是进一步了解AI4Science相关知识，然后使用深度学习的方法建模

一、基础知识点

（简单说明,详细见手册，链接见本节末）

1、Ai4Chemistry发展历史

在前期阶段，人们主要尝试使用不同的方法，尽可能地将化学知识和信息以计算机的形式进行存储，并以此为基础开始构建数据库；中期阶段，使用一些手动的特征工程对已有数据进行编码、特征提取等操作。再辅以传统的机器学习的方法，做一些预测；后期阶段，随着深度学习网络的广泛使用，这也导致后来又非常多的新型的分子指纹出现。

2、SMILES —— 最流行的将分子表示为序列类型数据的方法

SMILES，提出者Weininger et al[1]，全称是Simplified Molecular Input Line Entry System，是一种将化学分子用ASCII字符表示的方法，在化学信息学领域有着举足轻重的作用。

SMILES将化学分子中涉及的原子、键、电荷等信息，用对应的ASCII字符表示；环、侧链等化学结构信息，用特定的书写规范表达。

在SMILES中，原子由他们的化学符号表示，=表示双键、#表示三键、[]里面的内容表示侧基或者特殊原子（例如[Cu+2]表示带电+2电荷的Cu离子）。

3、分子指纹 —— 分子向量化

分子的指纹就像人的指纹一样，用于表示特定的分子。分子指纹是一个具有固定长度的位向量（即由0，1组成），其中，每个为1的值表示这个分子具有某些特定的化学结构。例如，对于一个只有长度为2的分子指纹，我们可以设定第一个值表示分子是否有甲基，第二个位置的值表示分子是都有苯环，那么[0,1]的分子指纹表示的就是一个有苯环而没有甲基的分子。

4、RDkit

RDkit是化学信息学中主要的工具，是开源的。网址：http://www.rdkit.org，支持WIN\MAC\Linux，可以被python、Java、C调用。

备注：在手册中也充分说明介绍了，在次就不进行阐述介绍了。手册链接：Datawhale (linklearner.com)

二、机器学习

机器学习按照目标可以分为分类任务（classification）和回归（regression）任务两大类。所谓分类任务，就是模型预测的结果是离散的值，例如类别；那么，回归任务中，模型预测的结果就是连续的值，例如房价等等。

就比说在task1的学习中，就使用了随机森林模型的机器学习方法。将多个决策树结合在一起，训练每个决策树的数据集都是随机有放回地从原数据中选出。预测的时候，输入会通过每个决策树进行预测，然后考虑每个树地输出结果，得到最终的预测值。

三、深度学习

深度学习可以归为机器学习的一个子集，主要通过神经网络学习数据的特征和分布。深度学习的一个重要进化是不再需要繁琐的特征工程，让神经网络自己从里面学习特征。

SMILES是一种以ASCII组成的序列，可以被理解为一种“化学语言”。既然是一种语言，那么很自然地想到了可以使用NLP中的方法对SMILES进行建模。

使用RNN对SMILES建模是早期的一个主要方法。RNN（Recurrent Neural Network）是处理序列数据的一把好手。RNN的网络每层除了会有自己的输出以外，还会输出一个隐向量到下一层。

图2 RNN的架构示意图

其中，每一层相当于做了一次线性变换：

h_n = \sigma(W_{hh}h_{n-1} + W_{hx}x_n + b_n)

每层的输出： y_n = Softmax(Vh_n + c)

通过隐向量的不断传递，序列后面的部分就通过“阅读”隐向量，获取前面序列的信息，从而提升学习能力。

但是RNN也有缺点：如果序列太长，那么两个相距比较远的字符之间的联系需要通过多个隐藏向量。这就像人和人之间传话一样，传递的人多了，很容易导致信息的损失或者扭曲。因此，它对长序列的记忆能力较弱。

同时，RNN需要一层一层地传递，所以并行能力差，比较容易出现梯度消失或梯度爆炸问题。

四、RNN模型训练解析

1、导入需要的库或包

import re
import time
import pandas as pd
from typing import List, Tuple
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, Subset

2、定义RNN模型

# 定义RNN模型
class RNNModel(nn.Module):
    def __init__(self, num_embed, input_size, hidden_size, output_size, num_layers, dropout, device):
        super(RNNModel, self).__init__()
        self.embed = nn.Embedding(num_embed, input_size) # 嵌入层，将词汇表中的单词转换为向量表示
        self.rnn = nn.RNN(input_size, hidden_size, num_layers=num_layers, 
                          batch_first=True, dropout=dropout, bidirectional=True) # RNN层，双向循环神经网络
        self.fc = nn.Sequential(nn.Linear(2 * num_layers * hidden_size, output_size), # 全连接层，用于输出预测结果
                                nn.Sigmoid(), # Sigmoid激活函数
                                nn.Linear(output_size, 1), # 再次全连接层，输出一个值
                                nn.Sigmoid()) # 再次应用Sigmoid激活函数

    def forward(self, x):
        # x : [bs, seq_len]
        x = self.embed(x) # 将输入序列中的每个元素转换为对应的向量表示
        # x : [bs, seq_len, input_size]
        _, hn = self.rnn(x) # 通过RNN层处理输入序列，返回最后一个时间步的隐藏状态
        hn = hn.transpose(0,1) # 转置隐藏状态张量，使其形状为[batch_size, num_layers*directions, hidden_size]
        z = hn.reshape(hn.shape[0], -1) # 将隐藏状态展平，以便输入到全连接层
        output = self.fc(z).squeeze(-1) # 通过全连接层得到最终输出，并压缩最后一个维度
        return output

初始化方法（__init__）：定义模型的结构，包括嵌入层、RNN层和全连接层。

self.embed: 嵌入层，将输入的索引转换为稠密向量。
self.rnn：双向RNN层，处理输入嵌入。
self.fc：全连接层，将RNN的输出转换为最终输出。

前向传播方法（forward）：定义数据通过模型的流动方式。

嵌入输入，得到嵌入向量。
通过RNN层，得到隐藏状态。
调整隐藏状态的形状并通过全连接层得到最终输出。

3、`SmilesTokenizer`：定义SMILES字符串的tokenizer

 数据处理部分
# tokenizer，鉴于SMILES的特性，这里需要自己定义tokenizer和vocab
# 这里直接将smiles str按字符拆分，并替换为词汇表中的序号
class Smiles_tokenizer():
    def __init__(self, pad_token, regex, vocab_file, max_length):
        self.pad_token = pad_token # 填充符号
        self.regex = regex # 正则表达式，用于匹配SMILES字符串中的元素
        self.vocab_file = vocab_file # 词汇表文件路径
        self.max_length = max_length # 最大序列长度

        with open(self.vocab_file, "r") as f:
            lines = f.readlines()
        lines = [line.strip("\n") for line in lines]
        vocab_dic = {}
        for index, token in enumerate(lines):
            vocab_dic[token] = index # 构建词汇表字典，键为元素，值为索引
        self.vocab_dic = vocab_dic

    def _regex_match(self, smiles):
        regex_string = r"(" + self.regex + r"|"
        regex_string += r".)"
        prog = re.compile(regex_string)

        tokenised = []
        for smi in smiles:
            tokens = prog.findall(smi) # 使用正则表达式匹配SMILES字符串中的元素
            if len(tokens) > self.max_length:
                tokens = tokens[:self.max_length] # 如果序列长度超过最大长度，截断
            tokenised.append(tokens) # 返回一个所有的字符串列表
        return tokenised
    
   def tokenize(self, smiles):
    # 使用正则表达式匹配SMILES字符串中的tokens
    tokens = self._regex_match(smiles)
    # 在每个token前后添加表示开始和结束的特殊标记：<CLS> 和 <SEP>
    tokens = [["<CLS>"] + token + ["<SEP>"] for token in tokens]
    # 对token序列进行填充，使其长度一致
    tokens = self._pad_seqs(tokens, self.pad_token)
    # 将token转换为对应的索引值
    token_idx = self._pad_token_to_idx(tokens)
    return tokens, token_idx

def _pad_seqs(self, seqs, pad_token):
    # 计算最长的序列长度
    pad_length = max([len(seq) for seq in seqs])
    # 对较短的序列进行填充，使其长度与最长序列相同
    padded = [seq + ([pad_token] * (pad_length - len(seq))) for seq in seqs]
    return padded

def _pad_token_to_idx(self, tokens):
    idx_list = []
    for token in tokens:
        tokens_idx = []
        for i in token:
            # 如果token已经在词汇表中，则获取其索引值；否则，将其添加到词汇表并分配一个新的索引值
            if i in self.vocab_dic.keys():
                tokens_idx.append(self.vocab_dic[i])
            else:
                self.vocab_dic[i] = max(self.vocab_dic.values()) + 1
                tokens_idx.append(self.vocab_dic[i])
        idx_list.append(tokens_idx)
    
    return idx_list

(1)初始化方法（__init__）：初始化标记器参数，读取词汇文件并创建词汇字典。self.vocab_dic: 词汇到索引的映射字典。

(2)正则表达式匹配（_regex_match）：使用正则表达式将SMILES字符串标记化prog.findall(smi): 使用正则表达式找到所有匹配的子串。

(3)标记化方法（tokenize）：将SMILES字符串标记化并填充到固定长度。

将标记转换为索引。
调用填充方法将标记填充到相同长度。_pad_seqs
添加和标记。<CLS><SEP>
self._regex_match(smiles): 调用正则表达式匹配方法。
(4)填充序列（_pad_seqs）：将序列填充到相同长度。
计算最长序列长度并填充短序列。

(5)标记转换为索引（_pad_token_to_idx）：将填充后的标记转换为索引。

4、读取数据并处理

# 定义一个函数，用于读取数据并处理
def read_data(file_path, train=True):
    # 使用pandas的read_csv方法读取CSV文件
    df = pd.read_csv(file_path)
    
    # 从DataFrame中提取各个列的数据，并将其转换为列表
    reactant1 = df["Reactant1"].tolist()
    reactant2 = df["Reactant2"].tolist()
    product = df["Product"].tolist()
    additive = df["Additive"].tolist()
    solvent = df["Solvent"].tolist()
    
    # 如果train为True，则提取Yield列的数据，否则创建一个与reactant1长度相同的全0列表
    if train:
        react_yield = df["Yield"].tolist()
    else:
        react_yield = [0 for i in range(len(reactant1))]
    
    # 将reactant拼到一起，之间用.分开。product也拼到一起，用>分开
    input_data_list = []
    for react1, react2, prod, addi, sol in zip(reactant1, reactant2, product, additive, solvent):
        input_info = ".".join([react1, react2])
        input_info = ">".join([input_info, prod])
        input_data_list.append(input_info)
    
    # 将输入数据和反应产量组合成一个元组列表
    output = [(react, y) for react, y in zip(input_data_list, react_yield)]

    '''
    input_data_list = []
    for react1, react2, prod, addi, sol in zip(reactant1, reactant2, product, additive, solvent):
        input_info = ".".join([react1, react2, addi, sol])
        input_info = ">".join([input_info, prod])
        input_data_list.append(input_info)
    output = [(react, y) for react, y in zip(input_data_list, react_yield)]
    '''

    # # 统计seq length，序列的长度是一个重要的参考，可以使用下面的代码统计查看以下序列长度的分布
    # seq_length = [len(i[0]) for i in output]
    # seq_length_400 = [len(i[0]) for i in output if len(i[0])>200]
    # print(len(seq_length_400) / len(seq_length))
    # seq_length.sort(reverse=True)
    # plt.plot(range(len(seq_length)), seq_length)
    # plt.title("templates frequence")
    # plt.show()
    
    # 返回处理后的数据
    return output

读取CSV文件并提取反应物、产物、添加剂和溶剂等列。
创建输入数据列表，将反应物和产物格式化为字符串。
如果是训练数据，读取产率，否则用0填充。
返回处理后的输入数据和产率

5、定义ReactionDataset

# 定义一个名为ReactionDataset的类，继承自Dataset类
class ReactionDataset(Dataset):
    # 初始化方法，接收一个包含字符串列表和浮点数的元组列表作为数据
    def __init__(self, data: List[Tuple[List[str], float]]):
        self.data = data  # 将传入的数据赋值给实例变量self.data
        
    # 重写Dataset类的__len__方法，返回数据集的长度
    def __len__(self):
        return len(self.data)  # 返回self.data的长度

    # 重写Dataset类的__getitem__方法，根据索引获取数据集中的元素
    def __getitem__(self, idx):
        return self.data[idx]  # 返回self.data中索引为idx的元素

自定义数据集类，用于加载反应数据。
__len__方法返回数据集的大小。
__getitem__方法根据索引返回对应的数据。

6、数据整理函数

def collate_fn(batch):
    # 定义一个正则表达式，用于匹配化学分子式中的特殊字符
    REGEX = r"\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9]"
    
    # 创建一个Smiles_tokenizer对象，用于将化学分子式转换为数字序列
    tokenizer = Smiles_tokenizer("<PAD>", REGEX, "../vocab_full.txt", max_length=300)
    
    # 初始化两个列表，分别用于存储化学分子式和对应的产量数据
    smi_list = []
    yield_list = []
    
    # 遍历输入的批次数据，提取化学分子式和产量数据
    for i in batch:
        smi_list.append(i[0])
        yield_list.append(i[1])
    
    # 使用tokenizer将化学分子式转换为数字序列，并将其转换为PyTorch张量
    tokenizer_batch = torch.tensor(tokenizer.tokenize(smi_list)[1])
    
    # 将产量数据转换为PyTorch张量
    yield_list = torch.tensor(yield_list)
    
    # 返回处理后的化学分子式数字序列和产量数据
    return tokenizer_batch, yield_list

定义了批处理函数，将数据批次转换为张量。
使用正则表达式初始化标记器。
标记化SMILES字符串并转换为张量。
将产率转换为张量并返回。

7、训练模型

def train():
    # 超参数设置
    # 数据集大小的一部分，int / int(len(dataset) * 1),或者你可以设置为数据集大小的一定比例，如 int(len(dataset) * 0.1)
    N = 10  
    NUM_EMBED = 294  # nn.Embedding()的嵌入维度
    INPUT_SIZE = 300  # 输入序列长度
    HIDDEN_SIZE = 512  # 隐藏层大小
    OUTPUT_SIZE = 512  # 输出层大小
    NUM_LAYERS = 10  # RNN层数
    DROPOUT = 0.2  # Dropout概率
    CLIP = 1  # 梯度裁剪阈值
    N_EPOCHS = 100  # 训练轮数
    LR = 0.0001  # 学习率

    start_time = time.time()  # 开始计时
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 选择设备（GPU或CPU）
    data = read_data("../dataset/round1_train_data.csv")  # 读取数据
    dataset = ReactionDataset(data)  # 创建数据集对象
    subset_indices = list(range(N))  # 获取数据集子集索引
    subset_dataset = Subset(dataset, subset_indices)  # 创建子集数据集
    train_loader = DataLoader(dataset, batch_size=128, shuffle=True, collate_fn=collate_fn)  # 创建数据加载器

    model = RNNModel(NUM_EMBED, INPUT_SIZE, HIDDEN_SIZE, OUTPUT_SIZE, NUM_LAYERS, DROPOUT, device).to(device)  # 创建模型并移动到指定设备
    model.train()  # 将模型设置为训练模式

    optimizer = optim.Adam(model.parameters(), lr=LR)  # 创建优化器
    criterion = nn.L1Loss()  # 使用平均绝对误差损失函数

    best_loss = 10  # 初始化最佳损失值
    for epoch in range(N_EPOCHS):  # 遍历每个训练轮次
        epoch_loss = 0  # 初始化当前轮次的损失值
        for i, (src, y) in enumerate(train_loader):  # 遍历数据加载器中的每个批次
            src, y = src.to(device), y.to(device)  # 将数据移动到指定设备
            optimizer.zero_grad()  # 清空梯度
            output = model(src)  # 前向传播计算输出
            loss = criterion(output, y)  # 计算损失
            loss.backward()  # 反向传播计算梯度
            torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP)  # 梯度裁剪
            optimizer.step()  # 更新模型参数
            epoch_loss += loss.item()  # 累加当前批次的损失值
        loss_in_a_epoch = epoch_loss / len(train_loader)  # 计算当前轮次的平均损失值
        print(f'Epoch: {epoch+1:02} | Train Loss: {loss_in_a_epoch:.3f}')  # 打印当前轮次的损失值
        if loss_in_a_epoch < best_loss:  # 如果当前轮次的损失值小于最佳损失值
            torch.save(model.state_dict(), '../model/RNN.pth')  # 保存模型状态字典
    end_time = time.time()  # 结束计时
    elapsed_time_minute = (end_time - start_time)/60  # 计算总运行时间（分钟）
    print(f"Total running time: {elapsed_time_minute:.2f} minutes")  # 打印总运行时间

if __name__ == '__main__':
    train()

训练参数: 定义训练所需的参数，如批量大小、学习率、训练轮数等。
设备设置：
读取数据: 调用read_data函数读取训练数据。
数据加载器: 创建数据加载器用于批量加载数据。
模型初始化：创建RNN模型实例并将其移动到设备。
优化器和损失函数：使用Adam优化器和L1损失函数。
训练循环: 迭代训练模型，计算损失，更新模型参数，保存最佳模型。

8、预测并生成提交文件

# 导入必要的库
import torch
from torch.utils.data import DataLoader
from dataset import ReactionDataset, collate_fn
from model import RNNModel
from data_utils import read_data

# 定义预测并生成提交文件的函数
def predicit_and_make_submit_file(model_file, output_file):
    # 设置模型参数
    NUM_EMBED = 294
    INPUT_SIZE = 300
    HIDDEN_SIZE = 512
    OUTPUT_SIZE = 512
    NUM_LAYERS = 10
    DROPOUT = 0.2
    
    # 检查是否有可用的GPU，如果有则使用GPU，否则使用CPU
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 读取测试数据
    test_data = read_data("../dataset/round1_test_data.csv", train=False)
    test_dataset = ReactionDataset(test_data)
    test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False, collate_fn=collate_fn) 

    # 初始化RNN模型并将其移动到相应的设备上
    model = RNNModel(NUM_EMBED, INPUT_SIZE, HIDDEN_SIZE, OUTPUT_SIZE, NUM_LAYERS, DROPOUT, device).to(device)
    
    # 加载最佳模型参数
    model.load_state_dict(torch.load(model_file))
    
    # 准备输出列表
    output_list = []
    
    # 遍历测试数据集并进行预测
    for i, (src, y) in enumerate(test_loader):
        src, y = src.to(device), y.to(device)
        with torch.no_grad():
            output = model(src)
            output_list += output.detach().tolist()
    
    # 准备提交文件的内容
    ans_str_lst = ['rxnid,Yield']
    for idx,y in enumerate(output_list):
        ans_str_lst.append(f'test{idx+1},{y:.4f}')
    
    # 将结果写入提交文件
    with open(output_file,'w') as fw:
        fw.writelines('\n'.join(ans_str_lst))

    print("done!!!")
    
# 调用函数进行预测并生成提交文件
predicit_and_make_submit_file("../model/RNN.pth", "../output/RNN_submit.txt")

设备设置：检查是否有GPU可用，设置预测设备。
读取数据: 调用read_data函数读取测试数据。
数据加载器: 创建数据加载器用于批量加载数据。
模型加载: 创建模型实例并加载训练好的模型参数。
预测循环: 迭代预测测试数据的产率并收集输出。
生成提交文件：将预测结果写入提交