【2025】Datawhale AI春訓營-RNA結構預測（AI+創新藥）-Task2筆記

本文對Task2提供的進階代碼進行理解。

任務描述

Task2的任務仍然是基于給定的RNA三維骨架結構，生成一個或多個RNA序列，使得這些序列能夠折疊并盡可能接近給定的目標三維骨架結構。這是一個RNA逆折疊的過程。

將RNA序列折疊成特定三維結構的過程是一個RNA折疊的過程。

在Task2中，繼續使用算法進行RNA逆折疊。評估標準是序列的恢復率，即算法生成的RNA序列在多大程度上能與真實能夠折疊成目標結構的RNA序列相似。

代碼理解

1、導入模塊

import os
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch_geometric
from torch_geometric.data import Data
from torch_geometric.nn import TransformerConv, LayerNorm
from torch_geometric.nn import radius_graph
from Bio import SeqIO
import math

2、配置參數

# 配置參數
class Config:seed = 42device = "cuda" if torch.cuda.is_available() else "cpu"batch_size = 16 if torch.cuda.is_available() else 8  # 根據顯存調整lr = 0.001epochs = 50seq_vocab = "AUCG"coord_dims = 7  hidden_dim = 256num_layers = 4  # 減少層數防止顯存溢出k_neighbors = 20  dropout = 0.1rbf_dim = 16num_heads = 4amp_enabled = True  # 混合精度訓練

3、定義幾何生成器

# 幾何特征生成器
class GeometricFeatures:@staticmethoddef rbf(D, D_min=0., D_max=20., D_count=16):device = D.deviceD_mu = torch.linspace(D_min, D_max, D_count, device=device)D_mu = D_mu.view(*[1]*len(D.shape), -1)D_sigma = (D_max - D_min) / D_countD_expand = D.unsqueeze(-1)return torch.exp(-((D_expand - D_mu)/D_sigma) ** 2)@staticmethoddef dihedrals(X, eps=1e-7):X = X.to(torch.float32)L = X.shape[0]dX = X[1:] - X[:-1]U = F.normalize(dX, dim=-1)# 計算連續三個向量u_prev = U[:-2]u_curr = U[1:-1]u_next = U[2:]# 計算法向量n_prev = F.normalize(torch.cross(u_prev, u_curr, dim=-1), dim=-1)n_curr = F.normalize(torch.cross(u_curr, u_next, dim=-1), dim=-1)# 計算二面角cosD = (n_prev * n_curr).sum(-1)cosD = torch.clamp(cosD, -1+eps, 1-eps)D = torch.sign((u_prev * n_curr).sum(-1)) * torch.acos(cosD)# 填充處理if D.shape[0] < L:D = F.pad(D, (0,0,0,L-D.shape[0]), "constant", 0)return torch.stack([torch.cos(D[:,:5]), torch.sin(D[:,:5])], -1).view(L,-1)@staticmethoddef direction_feature(X):dX = X[1:] - X[:-1]return F.pad(F.normalize(dX, dim=-1), (0,0,0,1))

4、定義圖構建器

# 圖構建器
class RNAGraphBuilder:@staticmethoddef build_graph(coord, seq):assert coord.shape[1:] == (7,3), f"坐標維度錯誤: {coord.shape}"coord = torch.tensor(coord, dtype=torch.float32)# 節點特征node_feats = [coord.view(-1, 7 * 3),  # [L,21]GeometricFeatures.dihedrals(coord[:,:6,:]),  # [L,10]GeometricFeatures.direction_feature(coord[:,4,:])  # [L,3]]x = torch.cat(node_feats, dim=-1)  # [L,34]# 邊構建pos = coord[:,4,:]edge_index = radius_graph(pos, r=20.0, max_num_neighbors=Config.k_neighbors)# 邊特征row, col = edge_indexedge_vec = pos[row] - pos[col]edge_dist = torch.norm(edge_vec, dim=-1, keepdim=True)edge_feat = torch.cat([GeometricFeatures.rbf(edge_dist).squeeze(1),  # [E,16]F.normalize(edge_vec, dim=-1)  # [E,3]], dim=-1)  # [E,19]# 標簽y = torch.tensor([Config.seq_vocab.index(c) for c in seq], dtype=torch.long)return Data(x=x, edge_index=edge_index, edge_attr=edge_feat, y=y)

5、定義模型結構

# 模型架構
class RNAGNN(nn.Module):def __init__(self):super().__init__()# 節點特征編碼self.feat_encoder = nn.Sequential(nn.Linear(34, Config.hidden_dim),nn.ReLU(),LayerNorm(Config.hidden_dim),nn.Dropout(Config.dropout))# 邊特征編碼（關鍵修復）self.edge_encoder = nn.Sequential(nn.Linear(19, Config.hidden_dim),nn.ReLU(),LayerNorm(Config.hidden_dim),nn.Dropout(Config.dropout))# Transformer卷積層self.convs = nn.ModuleList([TransformerConv(Config.hidden_dim,Config.hidden_dim // Config.num_heads,heads=Config.num_heads,edge_dim=Config.hidden_dim,  # 匹配編碼后維度dropout=Config.dropout) for _ in range(Config.num_layers)])# 殘差連接self.mlp_skip = nn.ModuleList([nn.Sequential(nn.Linear(Config.hidden_dim, Config.hidden_dim),nn.ReLU(),LayerNorm(Config.hidden_dim)) for _ in range(Config.num_layers)])# 分類頭self.cls_head = nn.Sequential(nn.Linear(Config.hidden_dim, Config.hidden_dim),nn.ReLU(),LayerNorm(Config.hidden_dim),nn.Dropout(Config.dropout),nn.Linear(Config.hidden_dim, len(Config.seq_vocab)))self.apply(self._init_weights)def _init_weights(self, module):if isinstance(module, nn.Linear):nn.init.xavier_uniform_(module.weight)if module.bias is not None:nn.init.constant_(module.bias, 0)def forward(self, data):x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr# 邊特征編碼（關鍵步驟）edge_attr = self.edge_encoder(edge_attr)  # [E,19] -> [E,256]# 節點編碼h = self.feat_encoder(x)# 消息傳遞for i, (conv, skip) in enumerate(zip(self.convs, self.mlp_skip)):h_res = conv(h, edge_index, edge_attr=edge_attr)h = h + skip(h_res)if i < len(self.convs)-1:h = F.relu(h)h = F.dropout(h, p=Config.dropout, training=self.training)return self.cls_head(h)

6、定義數據增強類

# 數據增強
class CoordTransform:@staticmethoddef random_rotation(coords):device = torch.device(Config.device)coords_tensor = torch.from_numpy(coords).float().to(device)angle = np.random.uniform(0, 2*math.pi)rot_mat = torch.tensor([[math.cos(angle), -math.sin(angle), 0],[math.sin(angle), math.cos(angle), 0],[0, 0, 1]], device=device)return (coords_tensor @ rot_mat.T).cpu().numpy()

7、定義數據集類

# 數據集類
class RNADataset(torch.utils.data.Dataset):def __init__(self, coords_dir, seqs_dir, augment=False):self.samples = []self.augment = augmentfor fname in os.listdir(coords_dir):# 加載坐標coord = np.load(os.path.join(coords_dir, fname))coord = np.nan_to_num(coord, nan=0.0)# 數據增強if self.augment and np.random.rand() > 0.5:coord = CoordTransform.random_rotation(coord)# 加載序列seq_id = os.path.splitext(fname)[0]seq_path = os.path.join(seqs_dir, f"{seq_id}.fasta")seq = str(next(SeqIO.parse(seq_path, "fasta")).seq)# 構建圖self.samples.append(RNAGraphBuilder.build_graph(coord, seq))def __len__(self): return len(self.samples)def __getitem__(self, idx): return self.samples[idx]

8、訓練函數

# 訓練函數
def train(model, loader, optimizer, scheduler, criterion):model.train()scaler = torch.cuda.amp.GradScaler(enabled=Config.amp_enabled)total_loss = 0for batch in loader:batch = batch.to(Config.device)optimizer.zero_grad()with torch.cuda.amp.autocast(enabled=Config.amp_enabled):logits = model(batch)loss = criterion(logits, batch.y)scaler.scale(loss).backward()torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)scaler.step(optimizer)scaler.update()total_loss += loss.item()scheduler.step()return total_loss / len(loader)

9、評估函數

# 評估函數
def evaluate(model, loader):model.eval()total_correct = total_nodes = 0with torch.no_grad():for batch in loader:batch = batch.to(Config.device)logits = model(batch)preds = logits.argmax(dim=1)total_correct += (preds == batch.y).sum().item()total_nodes += batch.y.size(0)return total_correct / total_nodes

10、主函數

if __name__ == "__main__":# 初始化torch.manual_seed(Config.seed)if torch.cuda.is_available():torch.cuda.manual_seed_all(Config.seed)torch.backends.cudnn.benchmark = True# 數據集train_set = RNADataset("./RNA_design_public/RNAdesignv1/train/coords","./RNA_design_public/RNAdesignv1/train/seqs",augment=True)# 劃分數據集train_size = int(0.8 * len(train_set))val_size = (len(train_set) - train_size) // 2test_size = len(train_set) - train_size - val_sizetrain_set, val_set, test_set = torch.utils.data.random_split(train_set, [train_size, val_size, test_size])# 數據加載train_loader = torch_geometric.loader.DataLoader(train_set, batch_size=Config.batch_size, shuffle=True,pin_memory=True,num_workers=4)val_loader = torch_geometric.loader.DataLoader(val_set, batch_size=Config.batch_size)test_loader = torch_geometric.loader.DataLoader(test_set, batch_size=Config.batch_size)# 模型初始化model = RNAGNN().to(Config.device)optimizer = optim.AdamW(model.parameters(), lr=Config.lr, weight_decay=0.01)scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=Config.epochs)criterion = nn.CrossEntropyLoss()# 訓練循環best_acc = 0for epoch in range(Config.epochs):train_loss = train(model, train_loader, optimizer, scheduler, criterion)val_acc = evaluate(model, val_loader)print(f"Epoch {epoch+1}/{Config.epochs} | Loss: {train_loss:.4f} | Val Acc: {val_acc:.4f}")if val_acc > best_acc:best_acc = val_acctorch.save(model.state_dict(), "best_model.pth")# 最終測試model.load_state_dict(torch.load("best_model.pth"))test_acc = evaluate(model, test_loader)print(f"\nFinal Test Accuracy: {test_acc:.4f}")