TPU (Tensor Processing Unit)
探索张量处理单元 (TPUs) 如何加速机器学习。学习优化用于 Edge TPU 的 Ultralytics YOLO26 以及云端训练以获得最大速度。
Tensor Processing Unit (TPU) 是一种专用应用特定集成电路 (ASIC),由 Google 专门设计用于加速 machine learning (ML) 工作负载。与处理广泛计算任务的通用处理器不同,TPU 从底层进行架构设计,旨在优化作为 neural networks 基础的大规模矩阵运算。这一特定重点使其能够实现极高的吞吐量和能效,使其成为现代 artificial intelligence (AI) 基础设施的基石,尤其是在 Google Cloud ecosystem 中。它们在减少训练复杂模型和大规模运行 real-time inference 所需时间方面发挥着至关重要的作用。
架构与功能#
TPU 的架构与传统处理器有很大不同。虽然标准的 CPU (Central Processing Unit) 擅长顺序任务和复杂逻辑,而 GPU (Graphics Processing Unit) 使用并行核心进行图形和通用计算,但 TPU 采用的是 systolic array architecture。这种设计使数据能够同时流经数千个乘法器,而无需为每次操作访问内存。通过最大化计算密度和最小化延迟,TPU 非常适合 deep learning (DL) 应用中繁重的线性代数运算。
这种专用硬件针对 TensorFlow 等框架进行了高度优化,并且日益得到 PyTorch 的支持,使开发者能够训练海量的 foundation models 或部署高效的边缘解决方案,而无需完全重写其代码库。
区分处理单元#
了解硬件环境对于优化 machine learning operations (MLOps) 至关重要。
- CPU: 计算机的通用“大脑”,非常适合顺序处理、数据预处理和处理复杂逻辑。它通常用于 data augmentation 管道,但对于繁重的矩阵数学运算速度较慢。
- GPU: 最初为图像渲染而构建的 GPU,由于其多功能性和大规模并行性,已成为 model training 的行业标准。它们非常适合训练像 Ultralytics YOLO26 这样灵活的模型。
- TPU: 一种专用加速器,用灵活性换取张量运算的原始速度。它旨在专门为神经网络计算最大化 FLOPS (floating-point operations per second),通常为特定的大规模工作负载提供出色的每瓦性能。
实际应用#
TPU 被部署在各种环境中,从海量的云集群到微小的边缘设备。
-
大语言模型训练: Google 利用称为 TPU Pod 的庞大互联集群来训练诸如 PaLM 和 Gemini 等巨大的 large language models (LLMs)。这些系统能够以传统硬件所需时间的一小部分处理 PB 级的 training data,从而加速 generative AI 的发展。
-
边缘 AI 与物联网: Coral Edge TPU 将这种加速带到了低功耗设备上。它支持高效的 computer vision (CV) 应用,例如在生产线上运行 object detection 以在本地识别缺陷。这允许在不依赖云连接的情况下进行即时决策,从而节省带宽和保护隐私。
在 Ultralytics 中使用 TPU#
开发者可以利用 TPU 加速来运行 Ultralytics 模型,特别是在使用 Ultralytics Platform 进行云端训练或导出用于边缘部署的模型时。例如,Edge TPU 要求模型专门为其架构进行量化和编译。
以下示例演示了如何将 Ultralytics YOLO26 模型导出为 TFLite 格式,这是针对 Edge TPU 进行编译之前的必要步骤:
from ultralytics import YOLO
# Load the latest lightweight YOLO26 nano model
model = YOLO("yolo26n.pt")
# Export the model to TFLite format
# This creates a '.tflite' file suitable for mobile and edge deployment
# Set int8=True for quantization, which is often required for Edge TPU performance
model.export(format="tflite", int8=True)导出后,可以使用 Edge TPU Compiler 进一步编译该模型,以便它能够在配有 Coral USB Accelerator 的 Raspberry Pi 等设备上高效运行。有关部署的更多详细信息,查阅 TFLite integration 文档会非常有帮助。






