DataWhale AI夏令营-催化反应速率预测(task2 笔记)

按照流程走

使用RNN网络建模SMILES序列

导包

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

定义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)
        self.fc = nn.Sequential(nn.Linear(2 * num_layers * hidden_size, output_size),
                                nn.Sigmoid(),
                                nn.Linear(output_size, 1),
                                nn.Sigmoid())

    def forward(self, x):
        # x : [bs, seq_len]
        x = self.embed(x)
        # x : [bs, seq_len, input_size]
        _, hn = self.rnn(x) # hn : [2*num_layers, bs, h_dim]
        hn = hn.transpose(0,1)
        z = hn.reshape(hn.shape[0], -1) # z shape: [bs, 2*num_layers*h_dim]
        output = self.fc(z).squeeze(-1) # output shape: [bs, 1]
        return output

__init__ 方法：

• num_embed：词汇表的大小，即嵌入层的词汇表大小。
• input_size：嵌入维度，即每个词的表示维度。
• hidden_size：RNN隐藏层的维度。
• output_size：RNN的输出层大小。
• num_layers：RNN的层数。
• dropout：dropout的概率，用于防止过拟合。
• device：用于指定计算设备（如GPU）

定义数据处理函数，以及tokenizer

将化学反应数据（以 SMILES 字符串形式表示）处理成可以输入到机器学习模型中的格式。处理流程包括分词、填充、索引化以及数据加载。

# import matplotlib.pyplot as plt
## 数据处理部分
# 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
        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)
            if len(tokens) > self.max_length:
                tokens = tokens[:self.max_length]
            tokenised.append(tokens) # 返回一个所有的字符串列表
        return tokenised
    
    def tokenize(self, smiles):
        tokens = self._regex_match(smiles)
        # 添加上表示开始和结束的token：<cls>, <end>
        tokens = [["<CLS>"] + token + ["<SEP>"] for token in tokens]
        tokens = self._pad_seqs(tokens, self.pad_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:
                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

# 读数据并处理
def read_data(file_path, train=True):
    df = pd.read_csv(file_path)
    reactant1 = df["Reactant1"].tolist()
    reactant2 = df["Reactant2"].tolist()
    product = df["Product"].tolist()
    additive = df["Additive"].tolist()
    solvent = df["Solvent"].tolist()
    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)]
    return output

class ReactionDataset(Dataset):
    def __init__(self, data: List[Tuple[List[str], float]]):
        self.data = data
        
    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]
    
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]"
    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_batch = torch.tensor(tokenizer.tokenize(smi_list)[1])
    yield_list = torch.tensor(yield_list)
    return tokenizer_batch, yield_list

训练模型

1. 超参数设置

N = 10  # 选择的数据子集的大小（这里是数据集的前10条数据）
NUM_EMBED = 294 # 嵌入层的维度
INPUT_SIZE = 300 # 输入的特征维度
HIDDEN_SIZE = 512 # RNN隐藏层的维度
OUTPUT_SIZE = 512 # 输出的维度
NUM_LAYERS = 10 # RNN的层数
DROPOUT = 0.2 # Dropout的比率
CLIP = 1 # 梯度裁剪的阈值
N_EPOCHS = 100 # 训练的总轮数
LR = 0.0001 # 学习率

2. 开始计时

start_time = time.time()  # 记录训练开始的时间

3. 选择设备

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

4. 读取数据

data = read_data("../dataset/round1_train_data.csv")
dataset = ReactionDataset(data)

5. 创建数据子集

subset_indices = list(range(N))
subset_dataset = Subset(dataset, subset_indices)

6. 数据加载器

train_loader = DataLoader(dataset, batch_size=128, shuffle=True, collate_fn=collate_fn)

7. 初始化模型

model = RNNModel(NUM_EMBED, INPUT_SIZE, HIDDEN_SIZE, OUTPUT_SIZE, NUM_LAYERS, DROPOUT, device).to(device)
model.train()

8. 优化器和损失函数

optimizer = optim.Adam(model.parameters(), lr=LR)
criterion = nn.L1Loss() # 使用 L1 损失（绝对误差）

9. 训练循环

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')

• 训练过程：
  • 对每一个 epoch（训练轮次），计算整个 epoch 的损失。
  • 遍历 train_loader 中的每一个批次。
  • 将数据转移到设备上（GPU/CPU）。
  • 进行前向传播计算输出。
  • 计算损失并进行反向传播。
  • 使用 clip_grad_norm_ 函数裁剪梯度，以防止梯度爆炸。
  • 更新模型参数。
  • 记录损失并计算平均损失。
• 模型保存：如果当前 epoch 的损失低于之前记录的最佳损失，则保存模型参数。

10. 结束计时

end_time = time.time()
elapsed_time_minute = (end_time - start_time) / 60
print(f"Total running time: {elapsed_time_minute:.2f} minutes")

生成结果

1. 加载训练好的模型。
2. 使用该模型对测试数据进行预测。
3. 将预测结果保存到一个指定格式的提交文件中。

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
    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) 

    model = RNNModel(NUM_EMBED, INPUT_SIZE, HIDDEN_SIZE, OUTPUT_SIZE, NUM_LAYERS, DROPOUT, device).to(device)
    # 加载最佳模型
    model.load_state_dict(torch.load(model_file))
    model.eval()
    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")

使用梯度提升回归模型

重新下载包

!pip install pandas
!pip install -U scikit-learn
!pip install rdkit

pip install scikit-learn==1.5.0

scikit-learn使用1.5.0版本

导入必要的库

import pickle
import pandas as pd
from tqdm import tqdm
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.model_selection import GridSearchCV
from rdkit.Chem import rdMolDescriptors
from rdkit import RDLogger, Chem
import numpy as np
RDLogger.DisableLog('rdApp.*')

pickle：用于保存和加载模型。

pandas：用于数据处理和读取。

tqdm：用于显示进度条。

sklearn.ensemble.GradientBoostingRegressor：梯度提升回归模型。

sklearn.model_selection.GridSearchCV：超参数调优工具。

rdkit.Chem.rdMolDescriptors 和 rdkit.Chem.Chem：用于计算分子指纹。

numpy：用于数值计算。

RDLogger：用于控制 RDKit 的日志信息。

定义分子指纹生成函数

def mfgen(mol, nBits=2048, radius=2):
    fp = rdMolDescriptors.GetMorganFingerprintAsBitVect(mol, radius=radius, nBits=nBits)
    return np.array(list(map(eval, list(fp.ToBitString()))))

mfgen：该函数生成一个分子的Morgan fingerprint（摩根指纹），并将其转换为一个NumPy数组。nBits指定位向量的长度，radius指定分子指纹的半径。

加载数据

def vec_cpd_lst(smi_lst):
    smi_set = list(set(smi_lst))
    smi_vec_map = {}
    for smi in tqdm(smi_set):
        mol = Chem.MolFromSmiles(smi)
        smi_vec_map[smi] = mfgen(mol)
    smi_vec_map[''] = np.zeros(2048)

    vec_lst = [smi_vec_map[smi] for smi in smi_lst]
    return np.array(vec_lst)

vec_cpd_lst：将SMILES（简化分子输入线性表示）转化为分子指纹，并为每个SMILES生成一个特征向量。进度条由tqdm显示。

 读取训练集和测试集

dataset_dir = '../dataset'   # 注：如果是在AI Studio上，将这里改为'dataset'

train_df = pd.read_csv(f'{dataset_dir}/round1_train_data.csv')
test_df = pd.read_csv(f'{dataset_dir}/round1_test_data.csv')

print(f'Training set size: {len(train_df)}, test set size: {len(test_df)}')

dataset_dir：数据集的目录。

pd.read_csv：从CSV文件中读取训练集和测试集数据。

打印训练集和测试集的大小。

处理训练集和测试集数据

train_rct1_smi = train_df['Reactant1'].to_list()
train_rct2_smi = train_df['Reactant2'].to_list()
train_add_smi = train_df['Additive'].to_list()
train_sol_smi = train_df['Solvent'].to_list()

train_rct1_fp = vec_cpd_lst(train_rct1_smi)
train_rct2_fp = vec_cpd_lst(train_rct2_smi)
train_add_fp = vec_cpd_lst(train_add_smi)
train_sol_fp = vec_cpd_lst(train_sol_smi)
train_x = np.concatenate([train_rct1_fp, train_rct2_fp, train_add_fp, train_sol_fp], axis=1)
train_y = train_df['Yield'].to_numpy()

test_rct1_smi = test_df['Reactant1'].to_list()
test_rct2_smi = test_df['Reactant2'].to_list()
test_add_smi = test_df['Additive'].to_list()
test_sol_smi = test_df['Solvent'].to_list()

test_rct1_fp = vec_cpd_lst(test_rct1_smi)
test_rct2_fp = vec_cpd_lst(test_rct2_smi)
test_add_fp = vec_cpd_lst(test_add_smi)
test_sol_fp = vec_cpd_lst(test_sol_smi)
test_x = np.concatenate([test_rct1_fp, test_rct2_fp, test_add_fp, test_sol_fp], axis=1)

从数据框中提取SMILES字符串。

将SMILES转换为分子指纹，并将所有指纹拼接成一个大的特征向量。

train_x 和 test_x 代表训练集和测试集的特征矩阵。

train_y 代表训练集的目标变量（产量）。

 定义和调优模型

param_grid = {
    'n_estimators': [300,],
    'learning_rate': [0.05,],
    'max_depth': [20,],
    'min_samples_split': [5, ],
    'min_samples_leaf': [2, ],
    'max_features': ['sqrt']
}

model = GradientBoostingRegressor()

grid_search = GridSearchCV(estimator=model, param_grid=param_grid, cv=5, n_jobs=-1, verbose=2)
grid_search.fit(train_x, train_y)

print("最优参数: ", grid_search.best_params_)

定义了超参数网格 param_grid 用于模型调优。

使用 GridSearchCV 进行超参数调优，通过交叉验证选择最佳参数。

打印最佳参数。

保存和加载模型

with open('./gradient_boosting_model.pkl', 'wb') as file:
    pickle.dump(best_model, file)

with open('gradient_boosting_model.pkl', 'rb') as file:
    loaded_model = pickle.load(file)

将训练好的最佳模型保存到文件中，并从文件中加载模型。

预测和输出结果

test_pred = loaded_model.predict(test_x)

ans_str_lst = ['rxnid,Yield']
for idx, y in enumerate(test_pred):
    ans_str_lst.append(f'test{idx+1},{y:.4f}')
with open('./submit.txt', 'w') as fw:
    fw.writelines('\n'.join(ans_str_lst))

使用加载的模型对测试集进行预测。

将预测结果格式化为字符串，并保存到 submit.txt 文件中。