计算生物——Code_Pytorch框架——蛋白靶点小分子药物对接亲和力预测_CPU（07.11-07.22）

写pytorch ，比我预期得更更更更费时.......

一.pytorch加载数据

1.  查看并保存数据集

使用数据：https://tdcommons.ai/multi_pred_tasks/dti/#bindingdb

利用数据集BindingDB进行回归任务，给定目标氨基酸序列/SMILES字符串，预测它们的结合亲和力

查看数据：

import requests
print(requests.__version__)
import tdc
print(dir(tdc))

from tdc.multi_pred import DTI
# 加载 DisGeNET 数据集
#data = GDA(name='DisGeNET')

data = DTI(name = 'BindingDB_Kd')

# 获取数据集的分割
split = data.get_split()

# 打印数据集信息
print("训练集样本数量:", len(split['train']))
print("验证集样本数量:", len(split['valid']))
print("测试集样本数量:", len(split['test']))

# 将训练集转换为DataFrame并展示前5行
train_df = pd.DataFrame(split['train'])
print(train_df.head())

# 将验证集和测试集也转换为DataFrame并展示前5行
valid_df = pd.DataFrame(split['valid'])
print(valid_df.head())

test_df = pd.DataFrame(split['test'])
print(test_df.head())

将数据集分为训练集，验证集，测试集

训练集展示：同一个target和不同的drug之间的不同的结合亲和力（Y）

将BindingDB输出在本地（分别为保存整体/保存train,valid,test）

# 将数据转换为 pandas DataFrame
combined_data = pd.concat([pd.DataFrame(split['train']),
                           pd.DataFrame(split['valid']),
                           pd.DataFrame(split['test'])])
# 将数据保存为 CSV 文件
combined_data.to_csv('BindingDB.csv', index=False)

print("Combined BindingDB dataset saved as BindingDB.csv.")

# 将数据保存为 CSV 文件
train_df.to_csv('BindingDB_train.csv', index=False)
valid_df.to_csv('BindingDB_valid.csv', index=False)
test_df.to_csv('BindingDB_test.csv', index=False)

print("BindingDB dataset saved as CSV files.")

2. 数据预处理

2.1 建立新环境并下载pytorch

先在VScode中新建立环境，并在环境内下好pytorch，详情见第一篇blog~

计算生物学习——Code_PyTorch(06.15-06.19)-CSDN博客

直到在bindingdb环境下，输入torch.cuda.is_available()，出现True才可以！

或者在VScode中的jupyter notebook中：

2.2 数据预处理

将 SMILES 字符串和氨基酸序列转换为模型可以处理的数值表示。

import numpy as np
import torch
print(torch.__version__)
from torch.utils.data import Dataset, DataLoader

import torch.nn as nn
import torch.nn.functional as F

import pandas as pd
from transformers import BertTokenizer, AutoTokenizer, BertModel, AutoModel 
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem import rdFingerprintGenerator

#ESM2 模型设置#
local_model_path = r'/home/embark/local_model'
# 加载本地的 tokenizer 和模型
tokenizer = AutoTokenizer.from_pretrained(local_model_path)
pytorch_path=r'/home/embark/local_model/pytorch_model.bin'
model = AutoModel.from_pretrained(local_model_path)

def sequence_to_embedding(seq):
    inputs = tokenizer(seq, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    return outputs.last_hidden_state.mean(dim=1).squeeze()

也就是说，对于氨基酸序列，用ESM2模型的加载方法还是依据下面博客中的第二种方法（确保能在软件里跑出来再使用！）

计算生物学习——Code_ESM-2输出蛋白质的特征表示_解决无法用代码线上加载模型（06.20-06.21）_esm2 蛋白模型-CSDN博客

对于SMILES，打算继续参考大佬loong_XL博客，用ChemBERTa(试过的其他方法，e.g.KPGT/MolBERT, 还没成功跑出来....)

ChemBERTa 化合物小分子的向量表示及相似检索-CSDN博客

怎么发现我没成功呢？（下面代码可以验证一下自己的函数设置有没有问题，能不能出来特征向量）

smiles = "NS(=O)(=O)c1ccc(S(=O)(=O)NCc2cccs2)s1"
sequence = "MSHHWGYGKHNGPEHWHKDFPIAKGERQSPVDIDTHTAKYDPSLKP"

smiles_vector = smiles_to_vector(smiles)
sequence_onehot = sequence_to_embedding(sequence)

print(smiles_vector)
print(sequence_onehot)

看到第一个smiles出来的tensor都是0，就很奇怪，一连换了几个smiles出来的都是0.，所有肯定SMILES的向量表示肯定没成功。

关于ChemBERTa，多更一个博客，详细点哈哈哈哈哈哈

计算生物——Code_SMILES的向量表示_ChemBERTa(07.16)-CSDN博客

#ChemBERTa 模型#
#路径是自己需要根据实际情况修改的哦#
local_model_path2 = r'/home/embark/chemBERTa/chemBERTa_files'
# 加载本地的 tokenizer 和模型
chemBERTa_tokenizer = RobertaTokenizer.from_pretrained(local_model_path2)
# 加载本地模型
chemBERTa_model = RobertaModel.from_pretrained(local_model_path2)

重新定义一下esm2模型

#ESM2 模型设置#
local_model_path = r'/home/embark/local_model'
# 加载本地的 tokenizer 和模型
esm2_tokenizer = AutoTokenizer.from_pretrained(local_model_path)

pytorch_path=r'/home/embark/local_model/pytorch_model.bin'
esm2_model = AutoModel.from_pretrained(local_model_path)

2.3 数据导入

第一步中保存的数据，直接放进VScode中，确认好路径。

import os
# 获取当前工作目录
current_dir = os.getcwd()
print(current_dir)
csv_file = os.path.join(current_dir, 'BindingDB_all.csv')
# 确认路径
print(f"CSV file path: {csv_file}")
# 读取 CSV 文件
dataa = pd.read_csv(csv_file)
data = dataa.dropna(subset=['Drug', 'Target', 'Y'])#删除这三列中有缺失值的行#
print(data.head())

随机抽取300个样本，是为了在CPU上训练快一点（GPU内存不够.......）

老师的建议是把两个的嵌入向量表示保存下来直接输入，这样不用再跑esm2和chemBERTa的模型，会快很多（俺还没试....）

# 随机抽取300个样本
data_sampled = data.sample(n=300, random_state=40)
# 对Y进行归一化处理
data_sampled['Y'] = (data_sampled['Y'] - data_sampled['Y'].min()) / (data_sampled['Y'].max() - data_sampled['Y'].min())
#print(data_sampled.head())

# 划分数据集为训练集、验证集和测试集
train_data, test_data = train_test_split(data_sampled, test_size=0.2, random_state=42)
train_data, val_data = train_test_split(train_data, test_size=0.25, random_state=42)  # 0.25 x 0.8 = 0.2
# 确保 train_data, val_data, 和 test_data 是 DataFrame
print(type(train_data))  # 应该输出 <class 'pandas.core.frame.DataFrame'>
print(type(val_data))    # 应该输出 <class 'pandas.core.frame.DataFrame'>
print(type(test_data))   # 应该输出 <class 'pandas.core.frame.DataFrame'>
print("train_data:", train_data)

# 计算Drug列中每一行序列的长度
data_sampled['Drug_length'] = data_sampled['Drug'].apply(len)
# 计算Target列中每一行序列的长度
data_sampled['Target_length'] = data_sampled['Target'].apply(len)
# 获取Drug列中序列长度的最大值
max_drug_length = data_sampled['Drug_length'].max()
# 获取Target列中序列长度的最大值
max_target_length = data_sampled['Target_length'].max()
print(f"Drug列中序列长度的最大值: {max_drug_length}")
print(f"Target列中序列长度的最大值: {max_target_length}")

运行结果如下，不难看到说，Drug列和Target列中都有长数据，这会使得每一行数据所生成的嵌入向量的形状不一致，有的很大，有的很小，需要长序列处理。（不然会出错...）

2.4 长序列处理

如果直接用数据集中的序列和SMILES，会出现：RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper_CUDA__index_select)

所以需要对长序列进行切片处理

接着上面ChemBERTa的博客，不难看到Output shape: torch.Size([1, 128, 384])，

简单的截断和填充效果并不好，甚至会使得长度超过128的序列信息被删除。（比如SMILES有300，那么后面的172个字母信息就没了）

使用以下代码在处理SMILES字符串时，模型输入被分成多个chunk，每个chunk的形状为torch.Size([1, 128])。这里的[1, 128]表示一个batch包含一个序列，该序列长度为128。经过所有chunk的处理后，最终得到的分子嵌入向量的形状为torch.Size([384])。这表示分子被嵌入到一个384维的向量空间中。

def process_chunks2(model, tokenizer, seq, max_length=128):
    # 将长序列分段
    tokens = tokenizer(seq, return_tensors="pt")['input_ids'].squeeze()
    chunks = [tokens[i:i + max_length] for i in range(0, len(tokens), max_length)]
    
    embeddings = []
    with torch.no_grad():
        for chunk in chunks:
            # 检查输入数据的有效性和类型
            if chunk.size(0) == 0:
                continue
            inputs = {'input_ids': chunk.unsqueeze(0).to(device_cpu)}
            print(f"Processing chunk with shape: {inputs['input_ids'].shape}")  # 打印输入数据的形状
            outputs = model(**inputs)
            chunk_embedding = outputs.last_hidden_state.mean(dim=1).squeeze()
            embeddings.append(chunk_embedding)
    
    # 合并所有段的嵌入
    final_embedding = torch.stack(embeddings).mean(dim=0)
    # 打印最终嵌入的形状
    print("Final embedding shape:", final_embedding.shape)
    return final_embedding

# 示例调用
seq = "NS(=O)(=O)c1ccc(S(=O)(=O)NCc2cccs2)s1" * 20  # 一个长序列
final_embedding = process_chunks2(model_, tokenizer_, seq, max_length=128)
print(final_embedding)

即便输入的序列不长，也可以使得最后得到的分子嵌入向量的形状为torch.Size([384])

现在可以确定这个函数可用了，即可使用代码：

# 定义处理长序列的函数
def process_chunks(model, tokenizer, seq, max_length=128):
    # 将长序列分段
    tokens = tokenizer(seq, return_tensors="pt", padding='max_length', truncation=True, max_length=max_length)['input_ids'].squeeze()
    chunks = [tokens[i:i + max_length] for i in range(0, len(tokens), max_length)]
    
    embeddings = []
    with torch.no_grad():
        for chunk in chunks:
            inputs = {'input_ids': chunk.unsqueeze(0)}
            outputs = model(**inputs)
            chunk_embedding = outputs.last_hidden_state.mean(dim=1).squeeze()
            embeddings.append(chunk_embedding)
    
    # 合并所有段的嵌入
    final_embedding = torch.stack(embeddings).mean(dim=0)
    return final_embedding

# 定义 smiles_to_vector 和 sequence_to_embedding 函数
def smiles_to_vector(seq):
    return process_chunks(chemBERTa_model, chemBERTa_tokenizer, seq, max_length=128)

def sequence_to_embedding(seq):
    return process_chunks(esm2_model, esm2_tokenizer, seq, max_length=128)

（一个函数，不管是对smiles还是sequence都能用）

3.定义数据集类

定义一个新的类TrainDataset，它继承自PyTorch的Dataset类。这个类用于创建一个自定义的数据集，可以与DataLoader一起使用来加载数据。

__init__函数，用于初始化类的实例变量。

• data：包含数据的DataFrame。
• smiles_to_vector：将SMILES字符串转换为向量表示的函数。
• sequence_to_embedding：将生物分子的序列转换为嵌入向量的函数。

__len__返回数据集的长度，即数据集中样本的数量。

__getitem__获取数据集中的单个样本。根据索引ind从DataFrame中提取一行数据，并分别获取药物的SMILES字符串、靶标的序列以及标签。

• smiles：药物的SMILES字符串。
• sequence：靶标的序列。
• label：标签，即药物与靶标的结合力。

使用assert语句检查转换后的向量和嵌入的形状是否符合预期。如果形状不匹配，则抛出异常并提供错误信息。

 #定义数据集类
class TrainDataset(Dataset):
    def __init__(self, data, smiles_to_vector, sequence_to_embedding):
        self.data = data
        self.smiles_to_vector = smiles_to_vector
        self.sequence_to_embedding = sequence_to_embedding

    def __len__(self):
        return len(self.data)

    def __getitem__(self, ind):
        row = self.data.iloc[ind]
        smiles = row['Drug']
        sequence = row['Target']
        label = row['Y']
        #使用提供的转换函数将SMILES字符串和序列转换为向量表示。
        smiles_vector = self.smiles_to_vector(smiles)
        sequence_embedding = self.sequence_to_embedding(sequence)
        # 检查形状是否符合要求
        assert smiles_vector.shape == torch.Size([384]), f"Unexpected shape {smiles_vector.shape} for drug {smiles}"
        assert sequence_embedding.shape == torch.Size([1280]), f"Unexpected shape {sequence_embedding.shape} for target {sequence}"
        
        return smiles_vector, sequence_embedding, torch.tensor(label, dtype=torch.float32)

查看数据集类是否定义好：

train_dataset = TrainDataset(train_data, smiles_to_vector, sequence_to_embedding)
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
# 获取一个批次的数据并打印
train_iter = iter(train_loader)
train_batch = next(train_iter)

print("Train batch:", train_batch)
print("Train batch shapes:", [t.shape for t in train_batch])

Train batch shapes: [torch.Size([16, 384]), torch.Size([16, 1280]), torch.Size([16])]

每个批次包含16个样本。这些形状与定义的smiles_to_vector和sequence_to_embedding函数的输出以及标签的数据类型和预期形状是一致的，说明数据加载和预处理工作正常。

• torch.Size([16, 384]):

  • 药物的SMILES向量的形状。
  • 批次大小为16，每个SMILES向量的维度为384。
  • 这意味着在一个批次中，有16个药物样本，每个样本被表示为一个384维的向量。

• torch.Size([16, 1280]):

  • 生物分子靶标序列的嵌入向量的形状。
  • 批次大小为16，每个序列嵌入的维度为1280。
  • 这意味着在一个批次中，有16个生物分子样本，每个样本被表示为一个1280维的向量。

• torch.Size([16]):

  • 标签的形状。
  • 批次大小为16，每个标签是一个标量值。
  • 这意味着在一个批次中，有16个标签，每个标签对应一个药物与生物分子靶标的结合力值。

准备好测试集

test_dataset3 = TrainDataset(test_data, smiles_to_vector, sequence_to_embedding)
test_loader = DataLoader(test_dataset3, batch_size=20, shuffle=False)

二.Pytorch 定义模型

因为我GPU用不了，所以只用CPU跑一下了，GPU能用就把第二行代码注释掉，第一行代码去掉#

#device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = torch.device('cpu')
# 定义一个简单的MLP模型
class MLP(nn.Module):
    def __init__(self, input_dim):
        super(MLP, self).__init__()
        self.linear1 = nn.Linear(input_dim, 512)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.2)
        self.linear2 = nn.Linear(512, 256)
        self.linear3 = nn.Linear(256, 1)

    def forward(self, drug, protein):
        x = torch.cat((drug, protein), dim=1)
        out = self.linear1(x)
        out = self.relu(out)
        out = self.dropout(out)
        out = self.linear2(out)
        out = self.relu(out)
        out = self.dropout(out)
        out = self.linear3(out)
        return out

# 计算更新后的输入维度
input_dim = 384 + 1280  # = 384 + 1280
model = MLP(input_dim).to(device)

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

print(model)

• __init__方法：定义网络的各层和激活函数。
  • self.linear1：输入层到第一个隐藏层，线性变换（全连接层），输入维度为input_dim，输出维度为512。
  • self.relu：ReLU激活函数。
  • self.dropout：Dropout层，用于防止过拟合，概率为0.2。
  • self.linear2：第二个隐藏层到第三个隐藏层，输入维度为512，输出维度为256。
  • self.linear3：第三个隐藏层到输出层，输入维度为256，输出维度为1。
• forward方法：定义前向传播过程。
  • 将药物和蛋白质向量拼接在一起（按维度1），然后通过一系列的线性变换、ReLU激活和Dropout，最后得到一个输出值。

input_dim：计算模型的输入维度，药物向量的维度是384，蛋白质向量的维度是1280，因此总的输入维度是384 + 1280。

model：实例化MLP模型，并将其移动到指定的设备（这里是CPU）。

criterion：定义损失函数为均方误差损失（MSELoss），用于回归问题。

optimizer：定义优化器为Adam，学习率设为0.001。

出现以下结果，则表示模型定义已成功

三. 训练模型

# 训练函数
def train_model(model, train_loader, criterion, optimizer, num_epochs):
    model.train()
    for epoch in range(num_epochs):
        running_loss = 0.0
        for i, (drug, protein, label) in enumerate(train_loader):
            optimizer.zero_grad()# 清空梯度
            outputs = model(drug, protein)
            loss = criterion(outputs.squeeze(), label)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}")

# 设置训练参数
num_epochs = 1
train_model(model, train_loader, criterion, optimizer, num_epochs)

这段代码定义了一个训练函数train_model，用于训练多层感知器（MLP）模型。训练过程中：

1. 每个epoch会遍历整个训练数据集。
2. 对每个batch，执行前向传播、计算损失、反向传播和参数更新。
3. 在每个epoch结束时，打印当前epoch的平均损失。

通过设置训练参数并调用该函数，模型会在1个epoch内进行训练，并输出损失信息。（为了时间原因，只能这样，如果追求更好的精度，可以num_epochs=100）

有以下结果：

（训练集有180个样本，16个为一批次，相当于有12个批次，耗时156min+）

四. 评估模型

model.eval()：将模型设置为评估模式。这一步确保在评估过程中禁用Dropout和BatchNorm层。由训练模型那儿可以预估一下，测试集有60个样本，以16个为一个批次，共四个批次，估计跑45-55min【CPU是真慢啊........】

from sklearn.metrics import r2_score, mean_squared_error
from scipy.stats import pearsonr

def evaluate_model(model, test_loader, criterion):
    model.eval()
    mse, r2, p_val, CI = 0, 0, 0, 0  # 初始化评估指标
    all_labels = []
    all_outputs = []
    
    print("Starting evaluation...")

    with torch.no_grad():
        for batch_idx, (drug, protein, label) in enumerate(test_loader):
            print(f"Processing batch {batch_idx+1}/{len(test_loader)}")
            outputs = model(drug, protein)
            loss = criterion(outputs.squeeze(), label)
            mse += loss.item()
            
            all_labels.append(label.cpu().numpy())
            all_outputs.append(outputs.squeeze().cpu().numpy())

    print("Calculating final metrics...") 
    
    mse /= len(test_loader)
    all_labels = np.concatenate(all_labels)
    all_outputs = np.concatenate(all_outputs)
    
    r2 = r2_score(all_labels, all_outputs)
    p_val = pearsonr(all_labels, all_outputs)[0]
    CI = concordance_index(all_labels, all_outputs)
    
    print(f"Mean Squared Error: {mse:.4f}")
    print(f"R^2: {r2:.4f}")
    print(f"Pearson Correlation: {p_val:.4f}")
    print(f"Concordance Index: {CI:.4f}")
    
    return mse, r2, p_val, CI

# 评估模型
mse, r2, p_val, CI = evaluate_model(model, test_loader, criterion)
print(f"R^2: {r2}, Pearson Correlation: {p_val}, Concordance Index: {CI}")

with torch.no_grad()：上下文管理器，禁用梯度计算，以节省内存并提高评估速度。

均方误差（MSE=0.0005）是衡量预测值与实际值之间平均误差的平方和。值越小，模型预测越准确。

决定系数  衡量的是模型对数据的拟合优度，取值范围从 -∞ 到 1。

• 1 表示完美拟合。
• 0 表示模型的预测能力与简单的平均模型相当。
• 负值表示模型比简单的平均模型还要差。

=-0.0696 表明模型的预测效果比简单的平均模型还要差。尽管 MSE 很小，但 R2R^2R2 的负值表明模型可能没有捕捉到数据中的重要趋势或模式。

Pearson Correlation=0.1620 表示模型预测值与实际值之间存在一定程度的正相关性，但相关性很弱。这表明模型在捕捉实际值的变化趋势方面表现不佳。

一致性指数（CI）衡量的是预测值和实际值之间的排序一致性，取值范围从 0.5 到 1。

• 1 表示完全一致的排序。
• 0.5 表示随机排序。
• 低于 0.5 表示预测效果比随机排序还差。

CI=0.4551 表示模型的预测排序与实际排序之间的一致性较差，接近于随机排序。这进一步说明模型在捕捉数据的顺序关系方面表现不佳。

（因为为了快一点检验pytorch框架是否可行，只用了300的样本并且模型训练num_epochs = 1，之后可以在GPU上跑快一点，用全部样本量50000+，并且num_epochs = 100，再试试）