RoMa: 一个用于在 PyTorch 中处理 3D 旋转的轻量级库。

RoMa（代表“Rotation Manipulation”）提供了 3D 旋转表示之间的可微映射、从欧几里得空间到旋转空间的映射，以及与旋转相关的各种实用工具。

它基于 PyTorch 实现，旨在为机器学习和基于梯度的优化提供一个易于使用且效率合理的工具箱。

文档

最新文档请访问：https://naver.github.io/roma/。

以下是 RoMa 的一些使用示例：

import torch
import roma

# 为了方便起见，支持任意数量的批量维度。
batch_shape = (2, 3)

# 旋转表示之间的转换
rotvec = torch.randn(batch_shape + (3,))
q = roma.rotvec_to_unitquat(rotvec)
R = roma.unitquat_to_rotmat(q)
Rbis = roma.rotvec_to_rotmat(rotvec)
euler_angles = roma.unitquat_to_euler('xyz', q, degrees=True)

# 从任意输入回归旋转：
# 对 3x3 矩阵进行特殊 Procrustes 正交化
R1 = roma.special_procrustes(torch.randn(batch_shape + (3, 3)))
# 从 6D 表示转换
R2 = roma.special_gramschmidt(torch.randn(batch_shape + (3, 2)))
# 从 4x4 对称矩阵的 10 个系数
q = roma.symmatrixvec_to_unitquat(torch.randn(batch_shape + (10,)))

# 旋转空间上的度量
R1, R2 = roma.random_rotmat(size=5), roma.random_rotmat(size=5)
theta = roma.utils.rotmat_geodesic_distance(R1, R2)
cos_theta = roma.utils.rotmat_cosine_angle(R1.transpose(-2, -1) @ R2)

# 四元数操作
q_identity = roma.quat_product(roma.quat_conjugation(q), q)

# 旋转向量之间的球面插值（最短路径）
rotvec0, rotvec1 = torch.randn(batch_shape + (3,)), torch.randn(batch_shape + (3,))
rotvec_interpolated = roma.rotvec_slerp(rotvec0, rotvec1, steps)

# 由旋转部分 R 和平移部分 t 组成的刚性变换 T
t = torch.randn(batch_shape + (3,))
T = roma.Rigid(R, t)
# 变换的组合与逆运算
identity = T @ T.inverse()
# 将结果转换为一批 4x4 齐次矩阵
M = identity.to_homogeneous()

安装

安装 RoMa 最简单的方式是使用 pip：

pip install roma

或者，也可以直接从源代码仓库安装最新版本的 RoMa：

pip install git+https://github.com/naver/roma

对于旧版 PyTorch（torch<1.8）**，我们建议安装 torch-batch-svd，以便在 CUDA GPU 上显著加速 special_procrustes 操作。您可以使用函数 roma.utils.is_torch_batch_svd_available() 来检查该模块是否已正确加载。**对于较新的 PyTorch 版本（torch>=1.8），则不再需要或使用 torch-batch-svd。

许可证

RoMa，版权所有 © 2020 NAVER Corp.，采用 3-Clause BSD 许可证（详见 license）。

部分代码改编自 SciPy。文档由 Sphinx 支持生成、分发并展示，其他材料也一并包含在内（详见 notice）。

贡献

如果您有任何建议，请在 GitHub 上提交问题。我们也欢迎 Pull 请求。我们的目标是保持 RoMa 的可靠性和可维护性，并可能根据具体情况接受您的贡献（无论是建议还是 Pull 请求），以实现这一目标。

通过为 RoMa 做出贡献，您即表示同意将您有权提交的贡献（无论是建议还是 Pull 请求）按照其 LICENSE 进行许可。

参考文献

如需深入了解旋转空间上的可微映射，请参阅：

• __Romain Brégier，《流形上的深度回归：3D 旋转案例研究》。2021 年国际 3D 视觉会议 (3DV)，2021 年。

请在您的出版物中引用此工作：

@inproceedings{bregier2021deepregression,
	title={Deep Regression on Manifolds: a {3D} Rotation Case Study},
	author={Br{\'e}gier, Romain},
	journal={2021 International Conference on 3D Vision (3DV)},
	year={2021}
}

RoMa 快速上手指南

RoMa (Rotation Manipulation) 是一个轻量级的 PyTorch 库，专为处理 3D 旋转设计。它提供了不同 3D 旋转表示法之间的可微分映射、从欧几里得空间到旋转空间的映射以及各种旋转相关工具，非常适合机器学习和基于梯度的优化任务。

环境准备

• 操作系统: Linux, macOS, Windows
• Python: 建议 Python 3.6+
• 核心依赖:
  • torch (PyTorch): 版本 >= 1.8 (推荐)。
  • 注: 若使用旧版 PyTorch (torch < 1.8) 且在 CUDA GPU 上运行，建议额外安装 torch-batch-svd 以显著提升 special_procrustes 的速度。新版 PyTorch 无需此依赖。

安装步骤

推荐使用 pip 进行安装：

pip install roma

如果需要安装最新的开发版本，可以直接从源码仓库安装：

pip install git+https://github.com/naver/roma

(国内用户若下载缓慢，可尝试添加 -i https://pypi.tuna.tsinghua.edu.cn/simple 参数使用清华镜像源)

基本使用

RoMa 支持任意数量的批次维度（batch dimensions），以下是核心功能的最简示例：

import torch
import roma

# 定义批次形状 (例如: 2x3 的批次)
batch_shape = (2, 3)

# 1. 旋转表示法转换
# 生成随机旋转向量 (rotvec)
rotvec = torch.randn(batch_shape + (3,))
# 转换为四元数 (unitquat)
q = roma.rotvec_to_unitquat(rotvec)
# 转换为旋转矩阵 (rotmat)
R = roma.unitquat_to_rotmat(q)
# 或者直接由旋转向量转为旋转矩阵
Rbis = roma.rotvec_to_rotmat(rotvec)
# 转换为欧拉角 (xyz 顺序，角度制)
euler_angles = roma.unitquat_to_euler('xyz', q, degrees=True)

# 2. 从任意输入回归旋转
# 对 3x3 矩阵进行特殊 Procrustes 正交化
R1 = roma.special_procrustes(torch.randn(batch_shape + (3, 3)))
# 从 6D 表示转换 (Gram-Schmidt 正交化)
R2 = roma.special_gramschmidt(torch.randn(batch_shape + (3, 2)))

# 3. 旋转空间度量
R_rand1, R_rand2 = roma.random_rotmat(size=5), roma.random_rotmat(size=5)
# 计算测地线距离
theta = roma.utils.rotmat_geodesic_distance(R_rand1, R_rand2)
# 计算余弦角度
cos_theta = roma.utils.rotmat_cosine_angle(R_rand1.transpose(-2, -1) @ R_rand2)

# 4. 刚体变换 (旋转 + 平移)
t = torch.randn(batch_shape + (3,))
T = roma.Rigid(R, t)
# 组合与逆变换: T * T_inv = 单位变换
identity = T @ T.inverse()
# 转换为 4x4 齐次矩阵批次
M = identity.to_homogeneous()