0. Baseline的基础解析:
0.1.赛题的基础解析
我们先来看看赛题的官方解释
RNA (核糖核酸) 在细胞生命活动中扮演着至关重要的角色,从基因表达调控到催化生化反应,都离不开 RNA 的参与。RNA 的功能很大程度上取决于其三维 (3D) 结构。理解 RNA 的结构与功能之间的关系,是生物学和生物技术领域的核心挑战之一。RNA 折叠 是指 RNA 序列自发形成特定三维结构的过程。而 RNA 逆折叠 则是一个更具挑战性的问题,即基于给定的RNA 三维骨架结构设计出能够折叠成这种结构的 RNA 序列。本次赛题的核心是 RNA 逆折叠问题,具体来说,是基于给定的 RNA 三维骨架结构,生成一个或多个 RNA 序列,使得这些序列能够折叠并尽可能接近给定的目标三维骨架结构。评估标准是序列的恢复率 (recovery rate),即算法生成的 RNA 序列,在多大程度上与真实能够折叠成目标结构的 RNA 序列相似。恢复率越高,表明算法性能越好。
这个次比赛的本质是逆推RNA的折叠步骤,简单来说就是我有一根电话线,随手一扔,他就会变随机生成一个形状,这个形状会直接决定它在细胞中扮演的角色,假设我们随机一扔,RNA折叠出了某个特定的复杂的形状(比如爱心、五角星)我们需要找出RNA的相关数据,让别人拿到这条“线”之后也可以扔出这个特定的形状。
比赛的目标简单可以概括为我们需要设计一个模型,针对给定的RNA的形状来倒推出可能实现这个RNA的序列,评分标准是我们推测出的人造序列和自然界真实的序列有多相似(也就是回复率)
文件的I/O格式如下: 本次比赛中,数据集的输入为 RNA 骨架原子 和 RNA 侧链原子,具体包括以下 7 个原子:
- RNA 骨架原子(6 个):
- P(磷酸基团)
- O5'(5' 氧原子)
- C5'(5' 碳原子)
- C4'(4' 碳原子)
- C3'(3' 碳原子)
- O3'(3' 氧原子)
- RNA 侧链原子(1 个):
- N1(嘧啶类碱基)或 N9(嘌呤类碱基) 数据以 numpy 数组 格式提供,每个原子的坐标为三维坐标(x, y, z)。如果原始数据中某个原子不存在,则该位置会以 NaN 填充。
输出格式为:
输出为 RNA 序列,以 FASTA 格式 文件给出。FASTA 文件包含两部分:
- 序列标识行:以 > 开头,后跟序列的唯一标识符。
- 序列行:RNA 序列的碱基(A, U, C, G)按顺序排列。
0.2 速通baseline:
这次我们用回我们的老朋友魔搭平台,这里的baseline已经调配得很好了,这里我们直接跑一下看看baseline的结果,整体的acc在0.3左右
我们再来看看代码,整体而言这个代码主要是基于图神经网络(GNN)进行RNA序列设计的模型。大体可以分为配置类、图构建器、数据集类、模型类、训练和评估函数。
我们一个一个来看,首先是配置类:
class Config:
seed = 42
device = "cuda" if torch.cuda.is_available() else "cpu"
batch_size = 32
lr = 0.01
epochs = 20
seq_vocab = "AUCG"
coord_dims = 7 # 7个骨架点
hidden_dim = 128
k_neighbors = 5 # 每个节点的近邻数
整体来说没有特别需要注意的地方 ,一个比较标准的设置,这里注意骨架点分别对应RNA核苷酸的7个关键原子:P, O5', C5', C4', C3', O3', O2'(磷酸骨架+糖环)
图构建类,整体上来说是基于RNA的结构进行的抽象,这个地方比较值得留意的是将RNA骨架点坐标建模为几何图,而非传统的用序列来表示(其实也是采用这类模型最大的亮点,RNA本身也具有集合结构),这里主要的实现是k近邻,给的数值是k=5,我这里的推测主要原因是RNA二级结构(如茎环)通常在5-10个核苷酸范围内形成。
class RNAGraphBuilder:
@staticmethod
def build_graph(coord, seq):
"""将坐标和序列转换为图结构"""
num_nodes = coord.shape[0]
# 节点特征:展平每个节点的7个骨架点坐标
x = torch.tensor(coord.reshape(num_nodes, -1), dtype=torch.float32) # [N, 7*3]
# 边构建:基于序列顺序的k近邻连接
edge_index = []
for i in range(num_nodes):
# 连接前k和后k个节点
neighbors = list(range(max(0, i-Config.k_neighbors), i)) + \
list(range(i+1, min(num_nodes, i+1+Config.k_neighbors)))
for j in neighbors:
edge_index.append([i, j])
edge_index.append([j, i]) # 双向连接
edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()
# 节点标签
y = torch.tensor([Config.seq_vocab.index(c) for c in seq], dtype=torch.long)
return Data(x=x, edge_index=edge_index, y=y, num_nodes=num_nodes)
数据集类,简单来说就是对数据集进行一些数据清洗,比如用0来替换原本的NaN,整体上面是标准的基于BioPyhon来读取FASTA的过程。没有太多特别的地方
class RNADataset(torch.utils.data.Dataset):
def __init__(self, coords_dir, seqs_dir):
self.samples = []
# 读取所有数据并转换为图
for fname in os.listdir(coords_dir):
# 加载坐标数据
coord = np.load(os.path.join(coords_dir, fname)) # [L, 7, 3]
coord = np.nan_to_num(coord, nan=0.0) # 新增行:将NaN替换为0
# 加载对应序列
seq_id = os.path.splitext(fname)[0]
seq = next(SeqIO.parse(os.path.join(seqs_dir, f"{seq_id}.fasta"), "fasta")).seq
# 转换为图结构
graph = RNAGraphBuilder.build_graph(coord, str(seq))
self.samples.append(graph)
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
return self.samples[idx]
接下来是本次任务的baseline模型,特征编码器设计是一个比较经典的低维集合特征映射到高维语义空间,主要化石为了增强特征的可分性,之后有两层的GCN,可以聚合两条的另据信息,这里主要的设计主要还是为了捕获RNA中长程的相互作用,分类头的设计相对比较简单,直接就是一个线性层,将节点特征直接映射到AUCG上。
class GNNModel(nn.Module):
def __init__(self):
super().__init__()
# 特征编码
self.encoder = nn.Sequential(
nn.Linear(7*3, Config.hidden_dim),
nn.ReLU()
)
# GNN层
self.conv1 = GCNConv(Config.hidden_dim, Config.hidden_dim)
self.conv2 = GCNConv(Config.hidden_dim, Config.hidden_dim)
# 分类头
self.cls_head = nn.Sequential(
nn.Linear(Config.hidden_dim, len(Config.seq_vocab))
)
def forward(self, data):
# 节点特征编码
x = self.encoder(data.x) # [N, hidden]
# 图卷积
x = self.conv1(x, data.edge_index)
x = torch.relu(x)
x = self.conv2(x, data.edge_index)
x = torch.relu(x)
# 节点分类
logits = self.cls_head(x) # [N, 4]
return logits
整体上来说代码说是一个相对简单的GNN结构模型,主要的缺点在三个方向
- 邻接关系:这里的k邻近不一定很好地反映整体的RNA的相互作用,RNA本身是在三维空间上作用的,需要考虑动态调整边的关系,同时邻居节点本身的重要性也未必都是相同的。
- 特征工程:这里插值主要是用的0填充,原子之间会有类型差异,比如P原子与O5'原子重要性不同但权重相同
- 模型本身:2层的GCN导致模型的感受野受限,没办法很好地匹配长程的碱基配对
1.BaseLine调优:
我们还是一个一个的来看,config这里我们加大了k邻近的数量,同时采用了动态批次
class Config:
batch_size = 16 if torch.cuda.is_available() else 8 # 动态批次
k_neighbors = 20 # 空间邻居数提升
dropout = 0.1 # 新增正则化参数
amp_enabled = True # 混合精度开关
几何特征生成器这边我们用了二面角特征计算,同时增加了方向向量生成,采用了RBF网络结构来编码距离,RBF核函数能够将数据从低维空间映射到高维空间,使得在高维空间中数据更易于线性可分。这个架构描述核苷酸构象会更加合适
class GeometricFeatures:
@staticmethod
def rbf(D, D_min=0., D_max=20., D_count=16):
device = D.device
D_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_count
D_expand = D.unsqueeze(-1)
return torch.exp(-((D_expand - D_mu)/D_sigma) ** 2)
@staticmethod
def 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)
@staticmethod
def direction_feature(X):
dX = X[1:] - X[:-1]
return F.pad(F.normalize(dX, dim=-1), (0,0,0,1))
图构建器这块我们用空间邻接代替了序列邻接,同时采用了复合边特征生成
class RNAGraphBuilder:
@staticmethod
def 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_index
edge_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)
模型设计比较多东西可以说,首先是采用4层Transformer卷积层替代原始GCN,通过多头注意力机制(4头)动态学习节点间的重要性权重,使模型能够区分关键相互作用(如氢键配对)与普通空间邻近关系。这主要是保证模型可以识别重要的邻居
其次是独立的边特征编码网络,将19维原始边特征(距离RBF+方向向量)映射到256维隐空间,使几何信息深度参与消息传递计算。
最后通过引入残差跳跃连接结构,每层输出通过线性变换后与输入相加,缓解梯度消失问题,允许构建更深的4层网络架构。
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)
训练函数这块引入了全局梯度裁剪,防止深层网络梯度爆炸,同时加上了余弦退火策略,实现了学习率的自动调节。
def train(model, loader, optimizer, scheduler, criterion):
model.train()
scaler = torch.cuda.amp.GradScaler(enabled=Config.amp_enabled)
total_loss = 0
for 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)
我们不妨来看看调优的结果,整体的acc上升了
2. 额外的调优尝试(尝试解决测评机分数较低的问题)
我打算尝试一波K折交叉验证,简单来说将数据分为k个子集(k折),每次用1个子集作为测试集,其余作为训练集,重复k次并取平均结果。这能更全面地反映模型在不同数据子集上的表现,减少单次划分的偏差。
我留意到给定数据集训练出的模型在测试集上的表现会比在测评机上的表现更好,所以希望通过k折交叉验证来提示模型泛化能力。这里主要的考虑还是希望通过k折折交叉验证来保证模型在面对一些比较少出现的RNA结构的时候依然有较为充足的泛化能力
以下是代码实现(初步),还在调试中
def kfold_train():
torch.manual_seed(Config.seed)
np.random.seed(Config.seed)
full_dataset = RNADataset(
"./RNA_design_public/RNAdesignv1/train/coords",
"./RNA_design_public/RNAdesignv1/train/seqs",
augment=True
)
kfold = KFold(n_splits=Config.k_folds, shuffle=True, random_state=Config.seed)
fold_metrics = {
'train_loss': [],
'val_acc': [],
'test_acc': []
}
for fold, (train_idx, val_idx) in enumerate(kfold.split(full_dataset)):
print(f"\n=== Fold {fold+1}/{Config.k_folds} ===")
test_size = len(val_idx) // 2
val_idx, test_idx = val_idx[:-test_size], val_idx[-test_size:]
train_sampler = SubsetRandomSampler(train_idx)
val_sampler = SubsetRandomSampler(val_idx)
test_sampler = SubsetRandomSampler(test_idx)
train_loader = DataLoader(
full_dataset,
batch_size=Config.batch_size,
sampler=train_sampler,
num_workers=4,
pin_memory=True
)
val_loader = DataLoader(
full_dataset,
batch_size=Config.batch_size,
sampler=val_sampler
)
test_loader = DataLoader(
full_dataset,
batch_size=Config.batch_size,
sampler=test_sampler
)
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 = 0
no_improve = 0
for epoch in range(Config.epochs):
train_loss = train(model, train_loader, optimizer, scheduler, criterion)
val_acc = evaluate(model, val_loader)
if val_acc > best_acc:
best_acc = val_acc
no_improve = 0
torch.save(model.state_dict(), f"fold{fold}_best.pth")
else:
no_improve += 1
if no_improve >= Config.early_stop_patience:
print(f"Early stopping at epoch {epoch}")
break
print(f"Epoch {epoch+1}: Loss={train_loss:.4f}, Val Acc={val_acc:.4f}")
model.load_state_dict(torch.load(f"fold{fold}_best.pth"))
test_acc = evaluate(model, test_loader)
fold_metrics['train_loss'].append(train_loss)
fold_metrics['val_acc'].append(best_acc)
fold_metrics['test_acc'].append(test_acc)
print(f"Fold {fold+1} Results:")
print(f"Best Val Acc: {best_acc:.4f}")
print(f"Test Acc: {test_acc:.4f}")
print("\n=== Final K-Fold Results ===")
for metric in fold_metrics:
values = fold_metrics[metric]
print(f"{metric}: {np.mean(values):.4f} ± {np.std(values):.4f}")
3.踩坑的一些记录:
3.1Torch相关
这次的踩坑主要是torch-cluster 安装上面,新的代码使用了torch-cluster,这里我发现直接按照torch-cluster非常容易出现死循环,一个一个排查之后,发现其实是torch-cluster的一个依赖torch-parse非常容易卡住,最终导致整个安装过程一直死循环
这里的解决方案是手动从网站上下载下轮子文件然后安装解决,要手动看看自己的cuda版本和torch版本
pip list
查看cuda版本
nvcc -v
这里找到对应的下载然后上传到云端就可以了
安装命令
pip install 你的文件名.whl
3.2 构建中的问题:
欢天喜地打包好交了一下一看,怎么0分?
看了一下log,大约说的是docker里面没有安装cluster(这里忘记截图了),回一直弹出Error Processing,导致代码没跑起来,恍然大悟这里docker文档没配置好,
这里必须修改一下你的docker,提前安装好这些东西,务必注意你自己的测试模型环境和测评机的docker环境不一定一样
RUN apt-get update && apt-get install -y \
build-essential \
python3-dev \
&& rm -rf /var/lib/apt/lists/*
ADD . /app
WORKDIR /app
RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple && \
pip config set global.extra-index-url "https://download.pytorch.org/whl/cu118 https://data.pyg.org/whl/torch-2.0.0+你的虚拟机的版本.html"
RUN pip install \
torch-scatter==你的docker \
torch-sparse==你的docker \
torch-cluster==你的docker \
torch-geometric==你的docker \
-f https://data.pyg.org/whl/torch-2.0.0+你的cuda版本.html
RUN pip install -r requirements.txt
CMD ["sh", "/app/run.sh"]
解决完毕了