Torchdyn: 深度学习与微分方程的完美结合

Torchdyn简介

Torchdyn是一个由DiffEqML团队维护的开源PyTorch库,专门致力于数值深度学习,包括微分方程、积分变换和数值方法。它为研究人员和实践者提供了一套完整的工具,用于探索连续深度学习模型的前沿。

Torchdyn的核心目标是将微分方程、深度学习和动力系统的研究融合在一起。近年来,这一领域的兴趣被一系列开创性工作重新点燃,如神经常微分方程(Neural ODEs)和神经随机微分方程(Neural SDEs)等。现代深度学习框架如PyTorch,加上计算资源的不断提升,使得神经网络的连续版本能够真正实现,为经典机器学习问题提供了全新的视角。

Torchdyn的主要特性

Torchdyn提供了丰富的功能和工具,使构建数值深度学习模型变得简单高效:

神经微分方程模型:提供了NeuralODE和NeuralSDE等核心类,可以轻松构建和训练基于微分方程的神经网络模型。
数值求解器:包含一套完整的ODE求解器和灵敏度方法,扩展了torchdiffeq提供的功能。
隐式模型:支持构建和训练各种隐式深度学习模型。
GPU加速:所有数值方法都兼容GPU,可以充分利用硬件加速。
与PyTorch集成:可以无缝集成到现有的PyTorch项目中。
教程和示例:提供了丰富的教程和应用示例,帮助用户快速上手。
可扩展性:设计为社区协作项目,欢迎贡献新的模型、方法和应用。

快速开始

使用Torchdyn构建神经微分方程模型非常简单。以下是一个基本示例:

from torchdyn.core import NeuralODE
import torch.nn as nn

# 定义神经网络作为ODE函数
f = nn.Sequential(
    nn.Conv2d(1, 32, 3),
    nn.Softplus(),
    nn.Conv2d(32, 1, 3)
)

# 创建NeuralODE模型
node = NeuralODE(f)

这样就创建了一个可训练的神经ODE模型。Torchdyn的设计理念是灵活性,你可以将Torchdyn的组件与任何PyTorch模块结合,构建复合模型。

安装

安装Torchdyn非常简单,可以通过pip直接安装稳定版:

pip install torchdyn

对于开发者,也可以使用poetry建立虚拟环境进行开发。

文档和教程

Torchdyn提供了全面的在线文档,包含详细的API参考和使用指南。此外,在项目的tutorials目录下还有许多自包含的快速入门示例和教程,专为实践者和研究人员设计。

Torchdyn logo

Torchdyn的应用领域

Torchdyn探索了可微分编程如何释放深度学习的效力,以加速多个科学领域的进展:

控制系统:利用神经微分方程进行系统建模和控制。
流体动力学:模拟和预测复杂流体系统的行为。
复杂动力系统预测:建立高精度的动力系统预测模型。
计算机视觉:应用连续深度学习模型解决图像分类、分割等任务。
自然语言处理:探索连续时间模型在序列建模中的应用。
金融建模:使用神经随机微分方程进行金融时间序列分析和风险建模。

通过提供集中、易于访问的模型模板、教程和应用笔记本集合,Torchdyn希望加速这一领域的研究,最终将神经微分方程和隐式模型确立为控制、系统识别和一般机器学习任务的有效工具。

技术细节

Torchdyn利用了现代PyTorch的最佳实践,并使用pytorch-lightning处理训练循环、日志记录和一般性管理。它还利用了dgl的图神经网络API来构建图神经ODE。Torchdyn提供了完整的ODE求解器和灵敏度方法套件,扩展了torchdiffeq提供的功能。此外,它还轻度依赖torchsde和torchcde来支持随机微分方程和控制微分方程的功能。

社区贡献

Torchdyn被设计为一个社区协作项目。我们欢迎所有与连续和隐式深度学习相关的教程、模型变体、数值方法和应用的贡献。虽然没有特定的风格要求,但我们遵循Jeremy Howard的许多编码理念。

我们使用poetry来管理依赖、创建虚拟Python环境和打包。贡献者可以参考文档来设置开发环境。

引用Torchdyn

如果你在研究或应用项目中使用了Torchdyn,请引用以下文章:

@article{politorchdyn,
  title={TorchDyn: Implicit Models and Neural Numerical Methods in PyTorch},
  author={Poli, Michael and Massaroli, Stefano and Yamashita, Atsushi and Asama, Hajime and Park, Jinkyoo and Ermon, Stefano}
}