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

安装环境

之前已经安装其它环境，这次只安装pandas

!pip install pandas

对数据处理进行初步尝试

导包

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

定义tokenizer和vocab，并将smiles str按字符拆分，并换为词汇表中的序号

1. 初始化方法 __init__

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

pad_token：填充用的标记。

regex：用于分词的正则表达式。

vocab_file：词汇表文件的路径，包含所有的词汇。

max_length：SMILES 字符串的最大长度。

读取 vocab_file 文件，创建一个词汇表字典 vocab_dic，将每个词汇映射到一个唯一的索引。

2. _regex_match 方法

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

使用给定的正则表达式 self.regex 对 SMILES 字符串进行分词。

分词后的结果是一个列表，其中每个 SMILES 字符串被拆分成一个词汇列表。

如果分词后的结果长度超过了 max_length，则截断到 max_length。

3. tokenize 方法

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

调用 _regex_match 方法进行分词。

在每个分词结果的开头添加 <CLS>（开始标记），在结尾添加 <SEP>（结束标记）。

调用 _pad_seqs 方法对分词结果进行填充，使得所有序列具有相同的长度。

调用 _pad_token_to_idx 方法将分词结果转换为索引。

返回填充后的分词结果和对应的索引。

4. _pad_seqs 方法

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

计算序列的最大长度。

对每个序列进行填充，使得所有序列的长度都等于最大长度。

填充使用的是 pad_token。

5. _pad_token_to_idx 方法

def _pad_token_to_idx(self, tokens):
    idx_list = []
    new_vocab = []
    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:
                new_vocab.append(i)
                self.vocab_dic[i] = max(self.vocab_dic.values()) + 1
                tokens_idx.append(self.vocab_dic[i])
        idx_list.append(tokens_idx)

    with open("../new_vocab_list.txt", "a") as f:
        for i in new_vocab:
            f.write(i)
            f.write("\n")

    return idx_list

将每个分词结果转换为词汇表中的索引。

如果遇到词汇表中不存在的词汇，将其添加到词汇表并分配一个新的索引。

将新词汇保存到 ../new_vocab_list.txt 文件中。

返回转换后的索引列表。

6. _save_vocab 方法

def _save_vocab(self, vocab_path):
    with open(vocab_path, "w") as f:
        for i in self.vocab_dic.keys():
            f.write(i)
            f.write("\n")
    print("update new vocab!")

将词汇表保存到指定路径的文件中。

打印 “update new vocab!” 表示词汇表已更新。

处理数据

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\additive\solvent拼到一起，之间用.分开。product也拼到一起，用>>分开
    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([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

读取数据：

使用 pd.read_csv 读取 CSV 文件中的内容，然后将相应的列转换为列表。列名包括：Reactant1、Reactant2、Product、Additive 和 Solvent。根据函数参数 train 的不同，还有一个 Yield 列。

处理反应物和产物：

如果 train 参数为 True，则提取 Yield 列作为反应产率；如果为 False，则用零填充。使用 zip 函数将反应物、产物、添加剂和溶剂对应起来，然后拼接它们的字符串表示。注释掉的那行代码示例了如何拼接所有化学成分，但当前版本只拼接了反应物。

拼接格式：

反应物以句点 . 作为分隔符，产物使用大于号 > 作为分隔符，形成最终的输入格式。

输出格式：

最后，返回一个包含输入数据和相应产率的元组的列表。

定义数据集

定义数据集
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", 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

处理化学反应数据并准备批处理数据

模型

class TransformerEncoderModel(nn.Module):
    def __init__(self, input_dim, d_model, num_heads, fnn_dim, num_layers, dropout):
        super().__init__()
        self.embedding = nn.Embedding(input_dim, d_model)
        self.layerNorm = nn.LayerNorm(d_model)
        self.encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, 
                                                        nhead=num_heads, 
                                                        dim_feedforward=fnn_dim,
                                                        dropout=dropout,
                                                        batch_first=True,
                                                        norm_first=True # pre-layernorm
                                                        )
        self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, 
                                                         num_layers=num_layers,
                                                         norm=self.layerNorm)
        self.dropout = nn.Dropout(dropout)
        self.lc = nn.Sequential(nn.Linear(d_model, 256),
                                nn.Sigmoid(),
                                nn.Linear(256, 96),
                                nn.Sigmoid(),
                                nn.Linear(96, 1))

    def forward(self, src):
        # src shape: [batch_size, src_len]
        embedded = self.dropout(self.embedding(src))
        # embedded shape: [batch_size, src_len, d_model]
        outputs = self.transformer_encoder(embedded)
        # outputs shape: [batch_size, src_len, d_model]

        # fisrt
        z = outputs[:,0,:]
        # z = torch.sum(outputs, dim=1)
        # print(z)
        # z shape: [bs, d_model]
        outputs = self.lc(z)
        # print(outputs)
        # outputs shape: [bs, 1]
        return outputs.squeeze(-1)

调整学习率

def adjust_learning_rate(optimizer, epoch, start_lr):
    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
    lr = start_lr * (0.1 ** (epoch // 3))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

这个函数每经过 3 个 epoch，就将学习率减少到原来的 0.1 倍

训练

def train():
    ## super param
    loss_in_a_epoch = 0
    N = 10#10  int / int(len(dataset) * 1)  # 或者你可以设置为数据集大小的一定比例，如 int(len(dataset) * 0.1)
    INPUT_DIM = 292 # src length
    D_MODEL = 512
    NUM_HEADS = 4
    FNN_DIM = 1024
    NUM_LAYERS = 4
    DROPOUT = 0.2
    CLIP = 1 # CLIP value
    N_EPOCHS = 40
    LR = 1e-4
    
    start_time = time.time()  # 开始计时
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # device = '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 = TransformerEncoderModel(INPUT_DIM, D_MODEL, NUM_HEADS, FNN_DIM, NUM_LAYERS, DROPOUT)
    model = model.to(device)
    model.train()
    
    optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=0.01)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10)
    criterion = nn.MSELoss()

    best_valid_loss = 10
    for epoch in range(N_EPOCHS):
        epoch_loss = 0
        # adjust_learning_rate(optimizer, epoch, LR) # 动态调整学习率
        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.detach().item()
            
            if i % 50 == 0:
                print(f'Step: {i} | Train Loss: {epoch_loss:.4f}')
                
        scheduler.step(loss_in_a_epoch)
        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_valid_loss:
            best_valid_loss = loss_in_a_epoch
            # 在训练循环结束后保存模型
            torch.save(model.state_dict(), '../model/transformer.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()         

参数说明

loss_in_a_epoch：

这是一个用于记录每个训练周期（epoch）中损失的变量，通常用来监控训练过程的性能。

N：

通常是指训练过程中使用的某种特定的批量大小、步数或者其他与训练相关的超参数，但在这个上下文中没有明确说明其具体用途。

INPUT_DIM：

输入数据的维度。对于文本数据，这可能是词嵌入的维度。这里是292，表示输入特征的数量或长度。

D_MODEL：

模型中每个层的隐藏状态维度，也称为模型维度。在变换器（Transformer）模型中，D_MODEL是每个注意力头和前馈网络的维度。

NUM_HEADS：

多头自注意力机制中的注意力头数量。多头机制允许模型从不同的子空间中学习信息，提高模型的表现力。

FNN_DIM：

前馈神经网络（Feed-Forward Network）中的隐藏层维度。在变换器模型中，前馈网络通常包含两个全连接层，FNN_DIM是第二层的维度。

NUM_LAYERS：

变换器模型中的层数。包括编码器和解码器中的层数。在编码器中，这个值表示堆叠的自注意力和前馈网络的层数。

DROPOUT：

Dropout正则化的比例。Dropout是一种防止过拟合的方法，它在训练时随机丢弃部分神经元的输出。DROPOUT为0.2表示20%的神经元会被丢弃。

CLIP：

梯度裁剪的阈值。为了防止梯度爆炸，通常会对梯度进行裁剪，以保证其在一个合理的范围内。CLIP为1表示梯度会被裁剪到最大为1。

N_EPOCHS：

训练的总轮数。模型将遍历训练数据集N_EPOCHS次，直到完成训练。

LR：

学习率（Learning Rate）。这是优化算法中控制权重更新步伐的参数。LR为1e-4（0.0001）表示每次更新权重时步伐的大小。

结果

# 生成结果文件
def predicit_and_make_submit_file(model_file, output_file):
    INPUT_DIM = 292 # src length
    D_MODEL = 512
    NUM_HEADS = 4
    FNN_DIM = 1024
    NUM_LAYERS = 4
    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=128, shuffle=False, collate_fn=collate_fn) 

    model = TransformerEncoderModel(INPUT_DIM, D_MODEL, NUM_HEADS, FNN_DIM, NUM_LAYERS, DROPOUT).to(device)
    # 加载最佳模型
    model.load_state_dict(torch.load(model_file))
    model.eval()
    output_list = []
    for i, (src, y) in enumerate(test_loader):
        src = src.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))
    
predicit_and_make_submit_file("../model/transformer.pth",
                              "../output/result.txt")