Whisper JAX 深度利用指南

Whisper JAX 深度利用指南

项目地点:https://gitcode.com/gh_mirrors/wh/whisper-jax
1. 项目介绍

Whisper JAX 是一个由 Sanchit Gandhi 创建的开源项目，它是 OpenAI 的 Whisper 模型在 JAX 平台上的实现。相比于原始的 PyTorch 版本，Whisper JAX 提供了高达 70 倍的速率提拔，使其成为目前最快的 Whisper 实现之一。项目兼容 CPU、GPU 和 TPU 环境，利用 JAX 的并行计算能力（如 pmap）举行高效处置处罚。
2. 项目快速启动

首先确保安装了 JAX 和相关的依赖库。你可以通过以下命令安装：

1. pip install flax jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_releases.html

复制代码

接下来，安装 Whisper JAX 包：

1. pip install --upgrade --no-deps --force-reinstall git+https://github.com/sanchit-gandhi/whisper-jax.git

复制代码

为了快速体验 Whisper JAX，可以创建并运行一个管道实例：

1. from whisper_jax import FlaxWhisperPipeline
3. # 初始化管道
4. pipeline = FlaxWhisperPipeline("openai/whisper-large-v2")
6. # 编译模型
7. pipeline.model.get_flax_module().init(key=jax.random.PRNGKey(0))

复制代码

请注意，JIT 编译可能会在初次调用时执行，之后的调用将更快。
3. 应用案例和最佳实践

数据并行性

要利用数据并行性，可以通过 FlaxWhisperPipline 类的 generate 方法来实现。这个方法已经封装了 pmap 函数，可以在多个设备上并行处置处罚输入：

1. inputs = ["Hello, how are you?"] * batch_size
2. transcriptions = pipeline.generate(inputs)

复制代码

自定义权重转换

如果你已经有 PyTorch 的 Whisper 模型权重，可以转换为 Flax 格式：

1. import jax.numpy as jnp
2. checkpoint_id = "sanchit-gandhi/whisper-small-hi"
3. model = FlaxWhisperForConditionalGeneration.from_pretrained(
4.     checkpoint_id, from_pt=True
5. )
7. # 将转换后的权重推送到 Hugging Face Hub
8. model.push_to_hub(checkpoint_id)
10. # 使用 Flax 权重加载管道
11. pipeline = FlaxWhisperPipeline(checkpoint_id, dtype=jnp.bfloat16, batch_size=16)

复制代码

4. 典型生态项目

Whisper JAX 项目是基于 T5x 代码库构建的，因此它可以利用 T5x 的模型激活和数据并行化技术。此外，由于它与 Hugging Face Transformers 兼容，可以轻松集成到各种 NLP 工具和工作流中。例如，结合 Hugging Face Accelerate 库，可以进一步优化在 GPU 或 TPU 上的性能。
其他基于 JAX 的相关生态项目包括但不限于：

• Optax: JAX 中的优化器库。
  
• Flax: JAX 上的神经网络库，用于构建可微分程序。
  
• Haiku: 另一个轻量级的 JAX 神经网络库。
  

通过这些工具，开辟者可以大概充分利用 JAX 的机动性和高性能特性，构建高效的 NLP 办理方案。
    whisper-jax JAX implementation of OpenAI's Whisper model for up to 70x speed-up on TPU.  

项目地点: https://gitcode.com/gh_mirrors/wh/whisper-jax   

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。