wavetorch

概述

该 Python 包为 PyTorch 提供了循环神经网络（RNN）模块，用于计算标量波动方程在时域中的解。本包中的代码是我们近期论文中所展示结果的基础，在该论文中，我们证明了当语音元音的录音通过一个经过训练的非均匀材料分布传播时，可以直接对其波形进行分类。

该包不仅提供了一个求解波动方程的数值框架，还允许通过 PyTorch 的自动微分（automatic differentiation）框架自动计算解的梯度。这种梯度计算等价于近年来在光子器件的逆向设计与优化中广受欢迎的伴随变量法（Adjoint Variable Method, AVM）。

更多详细信息和讨论请参见我们的论文：

• T. W. Hughes*, I. A. D. Williamson*, M. Minkov, and S. Fan, Wave physics as an analog recurrent neural network, Science Advances, vol. 5, no. 12, p. eaay6946, Dec. 2019

组件

本包中的机器学习示例围绕元音识别任务设计，使用了 James Hillenbrand 教授网站提供的原始音频录音数据集。然而，本包所提供的核心模块（如下所述）也可应用于其他涉及时序数据的学习或逆向设计任务。

wavetorch 包提供了多个独立模块，每个模块都是 torch.nn.Module 的子类。这些模块可以组合起来建模波动方程，或者（潜在地）作为构建其他网络的组件。

• WaveRNN - 一个封装器，包含一个或多个 WaveSource 模块、零个或多个 WaveProbe 模块，以及一个 WaveCell 模块。WaveRNN 模块是对各个组件的便捷封装，并负责对波动方程进行时间步进求解。如果没有探针（probe），WaveRNN 的输出是随时间变化的标量场分布；如果存在探针，则默认输出为探针处的值，但也可以覆盖此行为以输出整个场分布。
  • WaveCell - 实现标量波动方程的单个时间步。
    • WaveGeometry - 该模块的子类实现了 WaveCell 模块所使用的物理域参数化。尽管几何模块继承自 torch.nn.Module，但它没有 forward() 方法，仅用于向 WaveCell 模块提供材料密度的参数化表示。之所以继承 torch.nn.Module，是为了让可训练参数能被 PyTorch 正确识别。
  • WaveSource - 实现向标量波动方程中注入波的源。
  • WaveProbe - 实现在由 WaveGeometry 定义的域内特定位置测量波幅（或强度）的探针。

使用方法

波的传播

参见 study/propagate.py

透镜的优化与逆向设计

参见 study/optimize_lens.py

元音识别

要使用配置文件 study/example.yml 中指定的设置训练模型，请从仓库的顶层目录执行以下命令：

python ./study/vowel_train.py ./study/example.yml

配置文件 study/example.yml 中包含注释，说明了元音数据的处理方式、问题的物理设定以及训练过程的配置。

训练过程中，优化进度将打印到屏幕上。每个 epoch 结束时，当前模型状态以及所有先前 epoch 和交叉验证折（cross validation folds）的模型状态与性能历史将被保存到文件中。

警告：根据批次大小（batch size）、窗口长度（window length）和元音数据的采样率（sample rate）（这些均在 YAML 配置文件中指定），梯度计算可能需要大量内存。建议从较小的批次大小开始，根据机器的承受能力逐步增加。

元音识别结果汇总

可通过以下命令生成先前保存到磁盘的训练模型的摘要：

python ./study/vowel_summary.py <PATH_TO_MODEL>

在元音识别过程中显示场快照

可通过以下命令生成随机选取的元音样本的标量场分布快照：

python ./study/vowel_analyze.py fields <PATH_TO_MODEL> --times 1500 2500 3500 ...

显示元音波形的短时傅里叶变换（STFT）

可通过以下命令生成接收信号的短时傅里叶变换（Short-Time Fourier Transform, STFT）矩阵，其中行对应输入元音，列对应特定探针（与混淆矩阵的分布一致）：

python ./study/vowel_analyze.py stft <PATH_TO_MODEL>

依赖项

• pytorch
• scikit-learn
• scikit-image
• librosa
• seaborn
• matplotlib
• numpy
• yaml
• pandas

wavetorch 快速上手指南

环境准备

系统要求

• Python 3.7 或更高版本
• 支持 CUDA 的 GPU（可选，用于加速训练）

前置依赖
wavetorch 依赖以下 Python 包：

• pytorch
• scikit-learn
• scikit-image
• librosa
• seaborn
• matplotlib
• numpy
• PyYAML
• pandas

建议使用 Conda 或 pip 管理环境。若在国内，可使用清华源等镜像加速安装。

安装步骤

1. 克隆仓库：

git clone https://github.com/twhughes/wavetorch.git
cd wavetorch

2. 安装依赖（推荐使用国内镜像加速）：

# 使用 pip + 清华源
pip install torch scikit-learn scikit-image librosa seaborn matplotlib numpy PyYAML pandas -i https://pypi.tuna.tsinghua.edu.cn/simple

注意：PyTorch 安装请根据你的 CUDA 版本前往 https://pytorch.org/get-started/locally/ 获取合适命令，国内用户也可通过清华源安装。

3. （可选）将 wavetorch 添加为开发包（便于修改代码）：

pip install -e .

基本使用

示例：训练元音识别模型

1. 使用示例配置文件训练模型：

python ./study/vowel_train.py ./study/example.yml

该命令将基于 example.yml 中定义的数据处理、物理参数和训练设置进行训练。训练过程中会自动保存模型状态。

2. 查看训练结果摘要：

python ./study/vowel_summary.py <PATH_TO_MODEL>

3. 可视化波场快照（例如在时间步 1500、2500、3500）：

python ./study/vowel_analyze.py fields <PATH_TO_MODEL> --times 1500 2500 3500

提示：首次运行建议减小 example.yml 中的 batch_size 和 window_length，以避免显存不足。