AF3 squeeze_features函数解读

AlphaFold3 data_transforms 模块的 squeeze_features 函数的作用去除 蛋白质特征张量中不必要的单维度(singleton dimensions)和重复维度 ,以使其适配 AlphaFold3 预期的输入格式。

源代码:

复制代码
def squeeze_features(protein):
    """Remove singleton and repeated dimensions in protein features."""
    protein["aatype"] = torch.argmax(protein["aatype"], dim=-1)
    for k in [
        "domain_name",
        "msa",
        "num_alignments",
        "seq_length",
        "sequence",
        "superfamily",
        "deletion_matrix",
        "resolution",
        "between_segment_residues",
        "residue_index",
        "template_all_atom_mask",
    ]:
        if k in protein:
            final_dim = protein[k].shape[-1]
            if isinstance(final_dim, int) and final_dim == 1:
                if torch.is_tensor(protein[k]):
                    protein[k] = torch.squeeze(protein[k], dim=-1)
                else:
                    protein[k] = np.squeeze(protein[k], axis=-1)

    for k in ["seq_length", "num_alignments"]:
        if k in protein:
            protein[k] = protein[k][0]

    return protein

源码解读:

  • 该函数接收 protein(一个 包含蛋白质特征的字典)作为输入。
  • 主要任务:
    1. 将 one-hot aatype 转换为索引表示
    2. 移除 shape 为 (N, ..., 1) 的单维度
    3. 提取 seq_lengthnum_alignments 的实际数值
Step 1: 处理 aatype
复制代码
protein["aatype"] = torch.argmax(protein["aatype"], dim=-1)
  • 输入 aatype(氨基酸类型)通常是 one-hot 编码
  • 通过 torch.argmax(..., dim=-1) 获取 索引
  • 目的 :简化 aatype 的数据表示,使其直接存储氨基酸索引,而不是 one-hot 矩阵。
Step 2: 移除单维度
复制代码
for k in [
    "domain_name",
    "msa",
    "num_alignments",
    "seq_length",
    "sequence",
    "superfamily",
    "deletion_matrix",
    "resolution",
    "between_segment_residues",
    "residue_index",
    "template_all_atom_mask",
]:
    if k in protein:
        final_dim = protein[k].shape[-1]  # 获取最后一维的大小
        if isinstance(final_dim, int) and final_dim == 1:
            if torch.is_tensor(protein[k]):
                protein[k] = torch.squeeze(protein[k], dim=-1)  # 去掉单维度
            else:
                protein[k] = np.squeeze(protein[k], axis=-1)
  • 遍历多个 protein 特征字段,检查它们是否存在。
  • 如果最后一维 final_dim1 ,说明这个维度是无意义的单维度 ,需要去除:
    • 如果是 PyTorch 张量torch.Tensor),使用 torch.squeeze(dim=-1)
    • 如果是 NumPy 数组 ,使用 np.squeeze(axis=-1)
Step 3: 处理 seq_lengthnum_alignments
复制代码
for k in ["seq_length", "num_alignments"]:
    if k in protein:
        protein[k] = protein[k][0]

seq_lengthnum_alignments 可能是 列表或张量 ,但它们的数值其实是一个单独的整数,因此需要转换成 标量值

结论

1️⃣ 转换 aatype : 从 one-hot 编码 转换成 索引表示

2️⃣ 移除无用的单维度 : 让 msa, resolution, deletion_matrix 等数据符合 AlphaFold3 预期格式。

3️⃣ 转换 seq_lengthnum_alignments 为标量: 确保它们不会以张量形式存在,而是整数。

💡 最终作用:保证输入数据的维度符合 AlphaFold3 训练时的输入要求,提高数据处理效率。

相关推荐
catchadmin1 小时前
PHP 快速集成 ChatGPT 用 AI 让你的应用更聪明
人工智能·后端·chatgpt·php
编程武士4 小时前
从50ms到30ms:YOLOv10部署中图像预处理的性能优化实践
人工智能·python·yolo·性能优化
max5006005 小时前
基于Meta Llama的二语习得学习者行为预测计算模型
人工智能·算法·机器学习·分类·数据挖掘·llama
mCell5 小时前
长期以来我对 LLM 的误解
深度学习·llm·ollama
月疯6 小时前
OPENCV摄像头读取视频
人工智能·opencv·音视频
极客天成ScaleFlash6 小时前
极客天成让统一存储从云原生‘进化’到 AI 原生: 不是版本升级,而是基因重组
人工智能·云原生
王哥儿聊AI6 小时前
Lynx:新一代个性化视频生成模型,单图即可生成视频,重新定义身份一致性与视觉质量
人工智能·算法·安全·机器学习·音视频·软件工程
_pinnacle_6 小时前
打开神经网络的黑箱(三) 卷积神经网络(CNN)的模型逻辑
人工智能·神经网络·cnn·黑箱·卷积网络
Ada's6 小时前
深度学习在自动驾驶上应用(二)
人工智能·深度学习·自动驾驶
张较瘦_7 小时前
[论文阅读] 人工智能 + 软件工程 | 从“人工扒日志”到“AI自动诊断”:LogCoT框架的3大核心创新
论文阅读·人工智能·软件工程