CUDA 下的 PyTorch 软 DTW

PyTorch 的 soft-DTW 快速 CUDA 实现。 基于 pytorch-softdtw，但速度可提升至 100 倍！ forward() 和 backward() 步骤均使用 CUDA 实现。

我的实现部分受到 “时间序列数据中的模式发现方法及其 GPU 加速” 的启发，其中提出了一种基于对角线的贝尔曼递推实现。

入门

此代码依赖于 PyTorch 和 Numba。 只需将 soft_dtw_cuda.py 包含到你的项目中，即可开始使用！

你还可以运行附带的性能分析器/测试（已在 Python v3.6 上测试），查看可以获得的速度提升：

git clone https://github.com/Maghoumi/pytorch-softdtw-cuda
cd pytorch-softdtw-cuda
python soft_dtw_cuda.py

示例用法

脚本中已提供示例代码。以下是一个快速示例：

from soft_dtw_cuda import SoftDTW

# 创建序列
batch_size, len_x, len_y, dims = 8, 15, 12, 5
x = torch.rand((batch_size, len_x, dims), requires_grad=True)
y = torch.rand((batch_size, len_y, dims))
# 将张量传输到 GPU
x = x.cuda()
y = y.cuda()

# 创建“损失函数”对象
sdtw = SoftDTW(use_cuda=True, gamma=0.1)

# 计算损失值
loss = sdtw(x, y)  # 类似于任何 torch.nn.xyzLoss()

# 汇总并调用 backward()
loss.mean().backward()

演示项目

请查看 DeepNAG，我们的深度非对抗式手势生成器。 我们展示了，使用 soft DTW 训练的基于 RNN 的手势生成器，其性能可以超越使用 GAN 框架训练的相同生成器。

引用

如果你在研究中使用此代码，请引用以下出版物：

@phdthesis{maghoumi2020dissertation,
  title={{深度循环网络用于手势识别与合成}},
  author={Mehran Maghoumi},
  year={2020},
  school={佛罗里达州奥兰多中央佛罗里达大学}
}

@inproceedings{maghoumi2021deepnag,
  title={DeepNAG：深度非对抗式手势生成},
  author={Maghoumi, Mehran 和 Taranta, Eugene Matthew 以及 LaViola, Joseph},
  booktitle={第 26 届国际智能用户界面会议},
  pages={213--223},
  year={2021}
}

常见问题解答：

太棒了！我能做些什么来帮忙？

如果你觉得这个仓库有帮助，请考虑给它加星。同时别忘了感谢 pytorch-softdtw 的作者提供的 CPU 实现。

此外，也欢迎通过提升性能、解决现有限制等方式为本项目做出贡献。非常欢迎 PR！

是否支持剪枝？

是的！使用 bandwidth 参数指定 Sakoe-Chiba 剪枝带宽。

运行速度如何？

这取决于你的批大小和序列长度。序列越长、批越大，代码运行得越快。

以下是我在 Intel Core-i7 12700K 和 Titan RTX 上得到的结果：

针对 batch_size=128, seq_len_a=17, seq_len_b=15, dims=2 的 forward() + backward() 时间进行性能分析...
    CPU:      0.004228143487125635
    GPU:      0.0014472737908363341
    加速比：  2.9214537801325924

针对 batch_size=512, seq_len_a=64, seq_len_b=64, dims=2 的 forward() + backward() 时间进行性能分析...
    CPU:      0.023894597217440604
    GPU:      0.003414902277290821
    加速比：  6.997154025853163

针对 batch_size=512, seq_len_a=256, seq_len_b=256, dims=2 的 forward() + backward() 时间进行性能分析...
    CPU:      0.5894654761068523
    GPU:      0.0343648319132626
    加速比：  17.153160463425888

请注意，此代码仍有大量优化空间（例如使用共享内存等 CUDA 优化）。非常欢迎贡献和改进！

结果的准确性如何？

这取决于输入的长度。由于代码的顺序性，输入序列越长，累积的数值误差就越大。特别是在 backward() 调用中，对于范围在 [0, 1) 内的均匀随机输入，最终的导数张量中可能会出现高达 1e-3 的浮点误差。

soft_dtw_cuda.py 中包含的单元测试会将结果与 CPU 实现进行验证。

有哪些限制？

一些限制包括：

1. 同一批次中的所有序列必须具有相同的长度和特征数量。
2. 输入序列长度不能超过 1024（受 CUDA 最大线程块大小的限制）。 如果序列过长，代码会发出警告，并回退到 CPU 实现。
3. 如果输入较长（但仍小于 1024），可能会耗尽 CUDA 资源。详情见下文。

我遇到了 CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES 错误。怎么办？

这意味着你的序列太长，GPU 无法启动足够数量的线程。这与上述“限制”中的第 4 点有关。我不确定是否可以在 Numba 中查询 CUDA 设备，以判断根据所需线程数是否可以启动内核。在这种情况下，建议使用 CPU 实现。

许可证

本项目采用 MIT 许可证。

pytorch-softdtw-cuda 快速上手指南

pytorch-softdtw-cuda 是 PyTorch 的软动态时间规整（Soft-DTW）损失函数的高速 CUDA 实现。相比纯 CPU 版本，它在处理长序列和大批量数据时速度可提升高达 100 倍，且同时支持前向传播和反向传播的 GPU 加速。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux 或 Windows（需安装 NVIDIA 驱动及 CUDA Toolkit）。
• Python 版本：推荐 Python 3.6 及以上版本。
• 核心依赖：
  • PyTorch (需包含 CUDA 支持)
  • Numba (用于编译 CUDA 内核)

国内加速建议： 安装依赖时，推荐使用清华或阿里镜像源以提升下载速度：

pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install numba -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

该工具设计为单文件模块，无需通过 pip 安装复杂的包，只需将核心脚本引入项目即可。

1. 克隆仓库 将代码库克隆到本地：

   git clone https://github.com/Maghoumi/pytorch-softdtw-cuda.git
   cd pytorch-softdtw-cuda
   

2. 验证安装（可选） 运行自带的性能测试脚本，确认环境配置正确并查看 GPU 加速效果：

   python soft_dtw_cuda.py
   

3. 集成到项目 将目录下的 soft_dtw_cuda.py 文件复制到您自己的项目文件夹中，即可直接导入使用。

基本使用

以下是最简化的使用示例，展示了如何构建张量、初始化损失函数并进行反向传播。

import torch
from soft_dtw_cuda import SoftDTW

# 1. 创建输入序列
# 参数说明：batch_size=8, 序列长度 x=15, 序列长度 y=12, 特征维度=5
batch_size, len_x, len_y, dims = 8, 15, 12, 5
x = torch.rand((batch_size, len_x, dims), requires_grad=True)
y = torch.rand((batch_size, len_y, dims))

# 2. 将张量转移到 GPU (必须步骤)
x = x.cuda()
y = y.cuda()

# 3. 初始化 SoftDTW 损失对象
# use_cuda=True 启用 GPU 加速，gamma 为平滑参数
sdtw = SoftDTW(use_cuda=True, gamma=0.1)

# 4. 计算损失值 (用法同 torch.nn 中的标准 Loss)
loss = sdtw(x, y)

# 5. 聚合损失并执行反向传播
loss.mean().backward()

注意事项：

• 同一批次（Batch）中的所有序列必须具有相同的长度和特征维度。
• 单个序列长度不建议超过 1024，否则可能因 CUDA 资源限制自动回退到 CPU 模式或报错。