Geoopt: 一个强大的黎曼优化PyTorch扩展库

geoopt

Geoopt简介

Geoopt是一个基于PyTorch的黎曼优化库,它提供了一系列在黎曼流形上进行优化的工具和算法。黎曼优化是一种在约束空间(流形)上进行优化的方法,在许多机器学习任务中都有重要应用。Geoopt的目标是让研究人员和工程师能够方便地在PyTorch中使用黎曼优化技术。

Geoopt的主要特性

提供了多种常用黎曼流形的实现,如欧氏空间、Stiefel流形、球面等。
实现了在黎曼流形上的SGD、Adam等优化算法。
提供了ManifoldTensor和ManifoldParameter等特殊张量类型,方便在流形上进行操作。
支持与现有PyTorch模型的无缝集成。
实现了一些黎曼流形上的采样算法,如RSGLD、RHMC等。
提供了实验性的geoopt.layers模块,可以构建基于黎曼几何的神经网络层。

安装和使用

Geoopt可以通过pip轻松安装:

pip install git+https://github.com/geoopt/geoopt.git

安装后,可以像使用普通的PyTorch优化器一样使用Geoopt:

import geoopt

optimizer = geoopt.optim.RiemannianSGD(model.parameters(), lr=0.01)

Geoopt中的流形

Geoopt实现了多种常用的黎曼流形,包括:

Euclidean: 欧氏空间
Stiefel: Stiefel流形
Sphere: 球面
BirkhoffPolytope: 双随机矩阵流形
PoincareBall: 庞加莱球模型
Lorentz: 双曲模型
ProductManifold: 积流形构造器
SymmetricPositiveDefinite: 对称正定矩阵流形

以下是使用Sphere流形的示例:

import geoopt
import torch

sphere = geoopt.manifolds.Sphere()
x = torch.randn(10, 3)
x_proj = sphere.projx(x)  # 将x投影到球面上

Geoopt中的优化器

Geoopt提供了几种在黎曼流形上的优化器:

RiemannianSGD: 黎曼SGD优化器
RiemannianAdam: 黎曼Adam优化器
RiemannianLineSearch: 黎曼线搜索优化器

这些优化器可以像普通的PyTorch优化器一样使用。例如:

optimizer = geoopt.optim.RiemannianAdam(model.parameters(), lr=0.01)

for epoch in range(num_epochs):
    optimizer.zero_grad()
    loss = loss_fn(model(data), target)
    loss.backward()
    optimizer.step()

特殊张量类型

Geoopt提供了两种特殊的张量类型:

ManifoldTensor: 带有流形属性的张量
ManifoldParameter: 带有流形属性的参数

这些张量类型可以在定义模型参数时使用:

from geoopt import ManifoldParameter
import geoopt

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = ManifoldParameter(
            torch.randn(10, 10), 
            manifold=geoopt.manifolds.Stiefel()
        )