PyTorch使用Torchdyn实现连续时间神经网络的代码示例_Python

torchdyn概述

torchdyn是基于pytorch构建的专业库，专注于连续深度学习和隐式神经网络模型（如neural odes）的开发。该库具有以下核心特性：

支持深度不变性和深度可变性的ode模型
提供多种数值求解算法（如runge-kutta法，dormand-prince法）
与pytorch lightning框架的无缝集成，便于训练流程管理

本教程将以经典的moons数据集为例，展示neural odes在分类问题中的应用。

数据集构建

首先，我们使用torchdyn内置的数据集生成工具创建实验数据：

 from torchdyn.datasets import toydataset  
 import matplotlib.pyplot as plt  
   
 # 生成示例数据
 d = toydataset()  
 x, yn = d.generate(n_samples=512, noise=1e-1, dataset_type='moons')  
 # 可视化数据集
 colors = ['orange', 'blue']  
 fig, ax = plt.subplots(figsize=(3, 3))  
 for i in range(len(x)):  
     ax.scatter(x[i, 0], x[i, 1], s=1, color=colors[yn[i].int()])  
 plt.show()

数据预处理

将生成的数据转换为pytorch张量格式，并构建训练数据加载器。torchdyn支持cpu和gpu计算，可根据硬件环境灵活选择：

 import torch  
 import torch.utils.data as data  
   
 device = torch.device("cpu")  # 如果使用gpu则改为'cuda'
 x_train = torch.tensor(x).to(device)  
 y_train = torch.longtensor(yn.long()).to(device)  
 train = data.tensordataset(x_train, y_train)  
 trainloader = data.dataloader(train, batch_size=len(x), shuffle=true)

neural ode模型构建

neural odes的核心组件是向量场（vector field），它通过神经网络定义了数据在连续深度域中的演化规律。以下代码展示了向量场的基本实现：

 import torch.nn as nn  
   
 # 定义向量场f
 f = nn.sequential(  
     nn.linear(2, 16),  
     nn.tanh(),  
     nn.linear(16, 2)  
 )

接下来，我们使用torchdyn的

neuralode

类定义neural ode模型。这个类接收向量场和求解器设置作为输入。

 from torchdyn.core import neuralode  
   
 t_span = torch.linspace(0, 1, 5)  # 时间跨度
 model = neuralode(f, sensitivity='adjoint', solver='dopri5').to(device)

类来管理训练过程：

 import pytorch_lightning as pl  
   
 class learner(pl.lightningmodule):  
     def __init__(self, t_span: torch.tensor, model: nn.module):  
         super().__init__()  
         self.model, self.t_span = model, t_span  
     def forward(self, x):  
         return self.model(x)  
     def training_step(self, batch, batch_idx):  
         x, y = batch  
         t_eval, y_hat = self.model(x, self.t_span)  
         y_hat = y_hat[-1]  # 选择轨迹的最后一个点
         loss = nn.crossentropyloss()(y_hat, y)  
         return {'loss': loss}  
     def configure_optimizers(self):  
         return torch.optim.adam(self.model.parameters(), lr=0.01)  
     def train_dataloader(self):  
         return trainloader

最后训练模型：

 learn = learner(t_span, model)  
 trainer = pl.trainer(max_epochs=200)  
 trainer.fit(learn)

实验结果可视化

深度域轨迹分析

训练完成后，我们可以观察数据样本在深度域（即ode的时间维度）中的演化轨迹：

 t_eval, trajectory = model(x_train, t_span)  
 trajectory = trajectory.detach().cpu()  
   
 fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(10, 2))  
 for i in range(500):  
     ax0.plot(t_span, trajectory[:, i, 0], alpha=0.1, color=colors[int(yn[i])])  
     ax1.plot(t_span, trajectory[:, i, 1], alpha=0.1, color=colors[int(yn[i])])  
 ax0.set_title("维度 0")  
 ax1.set_title("维度 1")  
 plt.show()

向量场可视化

通过可视化学习得到的向量场，我们可以直观理解模型的动力学特性：

 x = torch.linspace(trajectory[:, :, 0].min(), trajectory[:, :, 0].max(), 50)  
 y = torch.linspace(trajectory[:, :, 1].min(), trajectory[:, :, 1].max(), 50)  
 x, y = torch.meshgrid(x, y)  
 z = torch.cat([x.reshape(-1, 1), y.reshape(-1, 1)], 1)  
 f_eval = model.vf(0, z.to(device)).cpu().detach()  
   
 fx, fy = f_eval[:, 0], f_eval[:, 1]  
 fx, fy = fx.reshape(50, 50), fy.reshape(50, 50)  
 fig, ax = plt.subplots(figsize=(4, 4))  
 ax.streamplot(x.numpy(), y.numpy(), fx.numpy(), fy.numpy(), color='black')  
 plt.show()