编译

在 SpeechBrain 中编译模型可以潜在地提高速度并减少内存需求。SpeechBrain 继承了 PyTorch 支持的编译方法，包括即时 (JIT) 编译器以及 PyTorch 版本 >=2.0 中引入的 torch.compile 方法。

使用 torch.compile 编译

torch.compile 功能随 PyTorch 版本 >=2.0 引入，旨在逐步取代 JIT。尽管此功能很有价值，但仍处于 beta 阶段，改进正在进行中。请查阅 PyTorch 文档 了解更多信息。

如何使用 torch.compile

在 SpeechBrain 中编译所有模块非常简单。您可以在运行训练 recipe 时使用命令行中的 --compile 标志来启用编译。例如

python train.py train.yaml --data_folder=your/data/folder --compile

这将自动编译 YAML 文件 modules 部分下声明的所有模块。

请注意，您可能需要正确配置额外的编译标志（例如，--compile_mode, --compile_using_fullgraph, --compile_using_dynamic_shape_tracing）以确保模型成功编译或实现最佳性能。要深入了解它们的作用，请参阅 PyTorch 文档 中的说明。

在某些情况下，您可能只想编译特定的模块。为此，可以使用 compile_module_keys 在 YAML 文件中添加要编译的模块键列表。例如

compile_module_keys: [encoder, decoder]

这将只编译编码器和解码器模型，这些模型应该在使用相应键之前在 YAML 文件中声明。

请记住使用 --compile 标志调用训练脚本。

注意事项：编译模型可能是一个复杂的过程，并且可能需要一些时间。此外，在某些情况下可能会失败。通过编译实现的加速高度依赖于所使用的系统和 GPU。例如，像 A100 这样的高端 GPU 倾向于产生更好的加速，而使用 V100 GPU 可能观察不到显著的改进。我们支持此功能，并希望 torch.compile 将随着时间的推移不断改进。

使用 JIT 编译

JIT 是 PyTorch 支持的第一个编译方法。值得注意的是，JIT 预计很快将被 torch.compile 取代。请查阅 PyTorch 文档 了解更多信息。

如何使用 JIT

要使用 JIT 在 SpeechBrain 中编译所有模块，请在运行训练 recipe 时使用命令行中的 --jit 标志

python train.py train.yaml --data_folder=your/data/folder --jit

这将自动编译 YAML 文件 modules 部分下声明的所有模块。

如果您只想编译特定的模块，可以使用 jit_module_keys 在 YAML 文件中添加要编译的模块键列表。例如

jit_module_keys: [encoder, decoder]

这将只编译编码器和解码器模型，前提是它们在 YAML 文件中使用指定的键声明。

请记住使用 --jit 标志调用训练脚本。

注意事项：JIT 对支持的语法有特定要求，许多流行的 Python 语法不被支持。因此，在设计模型时考虑到 JIT，请确保它满足成功编译所需的语法要求。此外，通过 JIT 编译实现的加速效果因模型类型而异。我们发现它对于自定义 RNN（例如 SpeechBrain 的 TIMIT/ASR/CTC 中使用的 Li-GRU）最有利。自定义 RNN 通常需要“for 循环”，这在 Python 中可能很慢。使用 JIT 编译可以在这种情况下提供显著加速。