🤗 Optimum TPU

Optimum TPU 提供了所有必要的机制，以利用和优化在 Google Cloud TPU 设备 上运行的 AI 工作负载。

该 API 提供与 Hugging Face Transformers 相同的用户体验，只需进行最少的更改即可实现推理性能。

训练支持正在进行中，敬请关注！ 🚀

安装

Optimum TPU 旨在最大限度地减少利用 Google Cloud TPU 加速器的摩擦。因此，我们提供了一个可通过 pip 安装的软件包，以确保每个人都能轻松上手。

使用 pip 运行 Cloud TPU

pip install optimum-tpu

在 Docker 容器中运行 Cloud TPU

PyTorch

export TPUVM_IMAGE_URL=us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla
export TPUVM_IMAGE_VERSION=8f1dcd5b03f993e4da5c20d17c77aff6a5f22d5455f8eb042d2e4b16ac460526
docker pull
docker run -ti --rm --privileged --network=host ${TPUVM_IMAGE_URL}@sha256:${TPUVM_IMAGE_VERSION} bash

然后，您可以通过上面的 pip 指令安装 optimum-tpu。