PyTorch/XLA 文档¶
torch_xla 是一个 Python 包，它实现了 XLA 作为 PyTorch 的后端。
熟悉 API 在 TPU 上创建和训练 PyTorch 模型，只需极少更改。 |
高性能 在数千个 TPU 核心上扩展训练作业，同时保持高 MFU。 |
成本效益高 TPU 硬件和 XLA 编译器针对成本效益高的训练和推理进行了优化。 |
入门指南¶
使用 pip 安装。
pip install torch torch_xla[tpu]
验证安装
python -c "import torch_xla; print(torch_xla.__version__)"
python -c "import torch; import torch_xla; print(torch.tensor(1.0, device='xla').device)"
教程¶
学习基础知识
TPU 上的分布式训练
高级技术
在 GPU 上训练
