AWS Trainium & Inferentia 文档

Inferentia 导出器

Hugging Face's logo
加入 Hugging Face 社区

并获得增强文档体验的访问权限

开始

Inferentia 导出器

您可以使用 🤗 Optimum 将 PyTorch 模型导出到 Neuron,以在 AWS Inferntia 1Inferentia 2 上运行推理。

导出函数

对于每一代 Inferentia 加速器,都有一个导出函数,export_neuron 适用于 INF1,export_neuronx 适用于 INF2,但您可以直接使用导出函数 export,它将根据环境选择合适的导出函数。

此外,您可以通过 validate_model_outputs 检查导出模型是否有效,该函数会将 Neuron 设备上编译模型的输出与 CPU 上 PyTorch 模型的输出进行比较。