【2025】Datawhale AI春训营-蛋白质预测（AI+生命科学）-Task2笔记

【2025】Datawhale AI春训营-蛋白质预测（AI+生命科学）-Task2笔记

本文对Task2使用的代码进行理解。

任务描述

Task2的任务仍然是通过对反应中包含的蛋白质残基信息，运用深度学习模型构建蛋白质3D结构的隐式模型，从而达成准确预测蛋白质内在无序区域（IDRs）的目的。任务的评价指标是实验真实结果和预测结果的F1 score。

代码理解

1、导入模块

import argparse
import math
import pickleimport torch
import torch.nn as nn
import torch.nn.functional as Ffrom tqdm import tqdm
from omegaconf import OmegaConf
from sklearn.metrics import f1_score
from torch.utils.data import Dataset, DataLoader
from torch.nn import TransformerEncoderLayer, TransformerEncoder

2、定义氨基酸类型

restypes = ['A', 'R', 'N', 'D', 'C','Q', 'E', 'G', 'H', 'I','L', 'K', 'M', 'F', 'P','S', 'T', 'W', 'Y', 'V'
]
unsure_restype = 'X'
unknown_restype = 'U'

3、定义数据集创建方法

def make_dataset(data_config, train_rate=0.7, valid_rate=0.2):data_path = data_config.data_pathwith open(data_path, 'rb') as f:data = pickle.load(f)total_number = len(data)train_sep = int(total_number * train_rate)valid_sep = int(total_number * (train_rate + valid_rate))train_data_dicts = data[:train_sep]valid_data_dicts = data[train_sep:valid_sep]test_data_dicts = data[valid_sep:]train_dataset = DisProtDataset(train_data_dicts)valid_dataset = DisProtDataset(valid_data_dicts)test_dataset = DisProtDataset(test_data_dicts)return train_dataset, valid_dataset, test_dataset

4、定义数据集

class DisProtDataset(Dataset):def __init__(self, dict_data):sequences = [d['sequence'] for d in dict_data]labels = [d['label'] for d in dict_data]assert len(sequences) == len(labels)self.sequences = sequencesself.labels = labelsself.residue_mapping = {'X':20}self.residue_mapping.update(dict(zip(restypes, range(len(restypes)))))def __len__(self):return len(self.sequences)def __getitem__(self, idx):sequence = torch.zeros(len(self.sequences[idx]), len(self.residue_mapping))for i, c in enumerate(self.sequences[idx]):if c not in restypes:c = 'X'sequence[i][self.residue_mapping[c]] = 1label = torch.tensor([int(c) for c in self.labels[idx]], dtype=torch.long)return sequence, label

5、定义位置编码类

class PositionalEncoding(nn.Module):def __init__(self, d_model, dropout=0.0, max_len=40):super().__init__()position = torch.arange(max_len).unsqueeze(1)div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))pe = torch.zeros(1, max_len, d_model)pe[0, :, 0::2] = torch.sin(position * div_term)pe[0, :, 1::2] = torch.cos(position * div_term)self.register_buffer("pe", pe)self.dropout = nn.Dropout(p=dropout)def forward(self, x):if len(x.shape) == 3:x = x + self.pe[:, : x.size(1)]elif len(x.shape) == 4:x = x + self.pe[:, :x.size(1), None, :]return self.dropout(x)

6、定义模型

class DisProtModel(nn.Module):def __init__(self, model_config):super().__init__()self.d_model = model_config.d_modelself.n_head = model_config.n_headself.n_layer = model_config.n_layerself.input_layer = nn.Linear(model_config.i_dim, self.d_model)self.position_embed = PositionalEncoding(self.d_model, max_len=20000)self.input_norm = nn.LayerNorm(self.d_model)self.dropout_in = nn.Dropout(p=0.1)encoder_layer = TransformerEncoderLayer(d_model=self.d_model,nhead=self.n_head,activation='gelu',batch_first=True)self.transformer = TransformerEncoder(encoder_layer, num_layers=self.n_layer)self.output_layer = nn.Sequential(nn.Linear(self.d_model, self.d_model),nn.GELU(),nn.Dropout(p=0.1),nn.Linear(self.d_model, model_config.o_dim))def forward(self, x):x = self.input_layer(x)x  = self.position_embed(x)x = self.input_norm(x)x = self.dropout_in(x)x = self.transformer(x)x = self.output_layer(x)return x

7、定义指标评估方法

def metric_fn(pred, gt):pred = pred.detach().cpu()gt = gt.detach().cpu()pred_labels = torch.argmax(pred, dim=-1).view(-1)gt_labels = gt.view(-1)score = f1_score(y_true=gt_labels, y_pred=pred_labels, average='micro')return score

8、定义主函数

if __name__ == '__main__':device = 'cuda' if torch.cuda.is_available() else 'cpu'parser = argparse.ArgumentParser('IDRs prediction')parser.add_argument('--config_path', default='./config.yaml')args = parser.parse_args()config = OmegaConf.load(args.config_path)train_dataset, valid_dataset, test_dataset = make_dataset(config.data)train_dataloader = DataLoader(dataset=train_dataset, **config.train.dataloader)valid_dataloader = DataLoader(dataset=valid_dataset, batch_size=1, shuffle=False)model = DisProtModel(config.model)model = model.to(device)optimizer = torch.optim.AdamW(model.parameters(),lr=config.train.optimizer.lr,weight_decay=config.train.optimizer.weight_decay)loss_fn = nn.CrossEntropyLoss()model.eval()metric = 0.with torch.no_grad():for sequence, label in valid_dataloader:sequence = sequence.to(device)label = label.to(device)pred = model(sequence)metric += metric_fn(pred, label)print("init f1_score:", metric / len(valid_dataloader))for epoch in range(config.train.epochs):# train loopprogress_bar = tqdm(train_dataloader,initial=0,desc=f"epoch:{epoch:03d}",)model.train()total_loss = 0.for sequence, label in progress_bar:sequence = sequence.to(device)label = label.to(device)pred = model(sequence)loss = loss_fn(pred.permute(0, 2, 1), label)progress_bar.set_postfix(loss=loss.item())total_loss += loss.item()optimizer.zero_grad()loss.backward()optimizer.step()avg_loss = total_loss / len(train_dataloader)# valid loopmodel.eval()metric = 0.with torch.no_grad():for sequence, label in valid_dataloader:sequence = sequence.to(device)label = label.to(device)pred = model(sequence)metric += metric_fn(pred, label)print(f"avg_training_loss: {avg_loss}, f1_score: {metric / len(valid_dataloader)}")# 保存当前 epoch 的模型save_path = f"model.pkl"torch.save(model.state_dict(), save_path)print(f"Model saved to {save_path}")