目录

AF3 identity_trans函数解读

AlphaFold3 rigid_utils 模块的 identity_trans 函数的功能是生成带有批次维度的全零平移向量张量。

源代码:

复制代码
@lru_cache(maxsize=None)
def identity_trans(
    batch_dims: Tuple[int], 
    dtype: Optional[torch.dtype] = None,
    device: Optional[torch.device] = None, 
    requires_grad: bool = True,
) -> torch.Tensor:
    trans = torch.zeros(
        (*batch_dims, 3), 
        dtype=dtype, 
        device=device, 
        requires_grad=requires_grad
    )
    return trans

源码解读:

1. 函数定义
复制代码
def identity_trans(
    batch_dims: Tuple[int], 
    dtype: Optional[torch.dtype] = None,
    device: Optional[torch.device] = None, 
    requires_grad: bool = True,
) -> torch.Tensor:

参数解析

  • batch_dims : 表示输入的批次维度,比如 (16,) 就创建 16 个零向量 (16, 3)

  • dtype : 数据类型,float32, float64 之类的。

  • device : 张量在哪个设备上,cpucuda

  • requires_grad : 是否需要梯度,默认为 True,适合训练场景。

👉 目标 :创建一个形状 [*, 3] 的零向量,代表初始平移向量 (0, 0, 0)

2. 创建张量
复制代码
trans = torch.zeros(
    (*batch_dims, 3), 
    dtype=dtype, 
    device=device, 
    requires_grad=requires_grad
)

解析逐项看

  • (*batch_dims, 3) :

    • batch_dims 展开成多个批次维度,比如 (16,) 展开后就是 16

    • 3 是每个向量的长度,表示 (x, y, z) 三个方向的位移。

    • 最终形状类似 torch.Size([16, 3])

  • dtype : 控制数据类型,比如 torch.float32

  • device : 控制张量生成在哪个设备,比如 cuda:0

  • requires_grad : 如果 True,这个张量就会参与梯度计算(适合训练用)。

3. 返回张量
复制代码
  return trans

最终返回 形状 [*, 3] 的全零平移向量张量。

4. 总结

identity_trans() 的核心功能:

生成初始位移向量 ------ 形状 [*, 3] 的零向量,代表 "不移动"

支持多批次输入 ------ batch_dims 灵活扩展支持多维数据,比如 [(8, 4, 3)]

缓存加速 ------ 重复调用相同参数时,不重复创建张量,直接用缓存结果。

支持梯度训练 ------ 默认开启 requires_grad=True,可以在训练时更新平移向量。

本文是转载文章,点击查看原文
如有侵权,请联系 xyy@jishuzhan.net 删除
相关推荐
Wnq100726 分钟前
DEEPSEEK 唤醒企业视频第二春
人工智能·嵌入式硬件·物联网·机器人·音视频·iot
蹦蹦跳跳真可爱5899 分钟前
Python----计算机视觉处理(Opencv:梯度处理:filiter2D算子,Sobel,Laplacian)
人工智能·python·opencv·计算机视觉
三桥君2 小时前
DeepSeek助力文案,智能音箱如何改变你的生活?
人工智能·生活·智能音箱·deepseek
蹦蹦跳跳真可爱5892 小时前
Python----计算机视觉处理(Opencv:绘制图像轮廓:寻找轮廓,findContours()函数)
人工智能·python·opencv·计算机视觉
天才测试猿2 小时前
Python常用自动化测试框架—Pytest
自动化测试·软件测试·python·功能测试·测试工具·测试用例·pytest
Ronin-Lotus2 小时前
深度学习篇---卷积网络结构
人工智能·python·深度学习·cnn
FreakStudio2 小时前
一文速通Python并行计算:04 Python多线程编程-多线程同步(上)—基于条件变量、事件和屏障
python·嵌入式·开发工具·面向对象·线程同步·代码阅读·电子diy
石硕页3 小时前
智能体是如何学习文档的内容的?【deepseek生成】
人工智能·rag·智能体
訾博ZiBo3 小时前
AI日报 - 2025年03月28日
人工智能
机械心3 小时前
自动驾驶VLA模型技术解析与模型设计
人工智能·机器学习·自动驾驶·vla·端到端自动驾驶