新手求助：如何用Equiformer模型训练原子坐标数据的得分函数？

问题解决说明 #

你遇到的「维度必须为2」报错，核心原因是直接传入了单样本的N×118独热特征，而Equiformer要求输入特征带批次维度（即(batch_size, num_atoms, feat_dim)）；同时高维独热特征直接作为输入也易引发维度兼容问题。下面是适配你的数据的极简训练脚本，包含独热特征的嵌入处理：

完整训练脚本 #

import torch
import torch.nn as nn
from equiformer_pytorch import Equiformer

# 1. 定义模型：先将独热原子特征嵌入到低维，再传入Equiformer
class ScoreModel(nn.Module):
    def __init__(self, num_elements=118, embed_dim=64):
        super().__init__()
        # 独热特征嵌入层：N×118 → N×64
        self.atom_embedding = nn.Linear(num_elements, embed_dim)
        # 初始化Equiformer：输入特征维度为embed_dim，输出为单得分值
        self.equiformer = Equiformer(
            dim=embed_dim,
            dim_head=32,
            heads=4,
            num_layers=3,
            num_degrees=4,
            attend_self=True,
            reduce_dim_out=True,
            out_dim=1  # 输出标量得分
        )

    def forward(self, coords, atom_one_hot):
        # 处理原子特征：独热编码转低维嵌入
        atom_feats = self.atom_embedding(atom_one_hot)  # shape: (B, N, embed_dim)
        # Equiformer前向传播：coords为(B, N, 3)，feats为(B, N, embed_dim)
        score = self.equiformer(coords, feats=atom_feats)
        return score.squeeze()  # 压缩维度为(B,)

# 2. 模拟训练数据（批次大小为2，每个样本含10个原子）
# 实际训练中需用DataLoader处理可变原子数，此处用固定长度简化
batch_coords = torch.randn(2, 10, 3)  # (batch_size, num_atoms, 3)
batch_one_hot = torch.zeros(2, 10, 118)
# 随机分配原子类型
for b in range(2):
    for n in range(10):
        elem_idx = torch.randint(0, 118, (1,))
        batch_one_hot[b, n, elem_idx] = 1.0

# 3. 初始化模型与优化器
model = ScoreModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 4. 按要求定义空逻辑损失函数
def dummy_loss(pred_score, target):
    return torch.tensor(0.0, requires_grad=True)

# 5. 训练循环示例
model.train()
for epoch in range(10):
    optimizer.zero_grad()
    # 前向传播
    pred_scores = model(batch_coords, batch_one_hot)
    # 模拟目标得分（实际训练替换为真实标签）
    target_scores = torch.randn(2)
    # 计算损失
    loss = dummy_loss(pred_scores, target_scores)
    # 反向传播与参数更新
    loss.backward()
    optimizer.step()
    print(f"Epoch {epoch+1}, Loss: {loss.item()}")

关键注意点 #

• 独热特征处理：通过nn.Linear将118维独热编码转换为低维嵌入，既解决维度兼容问题，也提升模型训练效率
• 批次维度：所有输入（坐标、原子特征）必须包含批次维度(batch_size, num_atoms, ...)，这是你之前报错的核心原因
• 可变原子数：实际训练中若每个样本原子数量不同，需用torch.nn.utils.rnn.pad_sequence或自定义DataLoader做padding处理

内容的提问来源于stack exchange，提问作者Connor Brown