torchsde

项目介绍

torchsde 是一个基于 PyTorch 的库，专注于可微分随机微分方程（SDE）求解器的实现。这个库的独特之处在于它支持 GPU 加速，并且能够高效地进行反向传播。随机微分方程是一类用于描述含有随机影响的动态系统的数学模型，torchsde 致力于提供强大而灵活的工具来解决这些方程。

安装要求

安装 torchsde 非常简单，用户只需通过 pip 命令进行安装：

pip install torchsde

需要注意的是，用户的 Python 环境版本需为 3.8 及以上，同时需要安装 PyTorch，版本至少为 1.6.0。

使用示例

快速上手

torchsde 提供了一些简单的使用示例，让用户能够快速了解如何使用该库解决 SDE。下面的代码展示了一个基本的 SDE 求解流程：

import torch
import torchsde

batch_size, state_size, brownian_size = 32, 3, 2
t_size = 20

class SDE(torch.nn.Module):
    noise_type = 'general'
    sde_type = 'ito'

    def __init__(self):
        super().__init__()
        self.mu = torch.nn.Linear(state_size, 
                                  state_size)
        self.sigma = torch.nn.Linear(state_size, 
                                     state_size * brownian_size)

    # Drift
    def f(self, t, y):
        return self.mu(y)  # 形状为 (batch_size, state_size)

    # Diffusion
    def g(self, t, y):
        return self.sigma(y).view(batch_size, 
                                  state_size, 
                                  brownian_size)

sde = SDE()
y0 = torch.full((batch_size, state_size), 0.1)
ts = torch.linspace(0, 1, t_size)
# 初始状态 y0，将 SDE 求解于时间区间 [ts[0], ts[-1]] 上。
# ys 的形状为 (t_size, batch_size, state_size)
ys = torchsde.sdeint(sde, y0, ts)

更多示例

除了基本的求解示例，torchsde 还提供了其他一些复杂的示例。这些示例帮助用户了解更多关于求解 SDE 时的细微之处，比如如何在求解器中固定随机性和选择不同的噪声类型。

1. 潜在 SDE：此示例中，torchsde 被用于学习潜在的随机微分方程模型。用户可以通过运行脚本来实现，例如：

   python -m examples.latent_sde --train-dir <TRAIN_DIR>
   

   这是一个将 SDE 拟合到数据上的过程，同时将其正则化为类似于 Ornstein-Uhlenbeck 过程的先验。这种模型可以视为一种变分自编码器（VAE），其中的先验和后验分布都是 SDE。

2. 作为 GAN 的神经 SDE：这个示例展示了如何使用 SDE 作为生成对抗网络（GAN）中的生成器，同时使用一个神经受控微分方程（Neural CDE）作为判别器。用户可以通过以下命令运行：

   python -m examples.sde_gan
   

项目背景

torchsde 诞生于一个研究项目，而非官方的 Google 产品。它的开发旨在为机器学习研究人员提供一个强大的工具，以便进行与 SDE 相关的研究和应用。如果在研究中使用了这个代码库，开发者建议在相关学术论文中进行引用。

总而言之，torchsde 提供了一个简单易用且功能全面的平台，用于解决随机微分方程，这对需要处理随机动态系统或进行相关模拟的研究人员来说，是一个理想的选择。