AITemplate首頁、文檔和下載- Meta 開源的Python 框架- 科技資訊

AITemplate（AIT）是一個Python 框架，它將深度神經網絡轉化為CUDA（NVIDIA GPU）/ HIP（AMD GPU）C++ 程式碼，以實現快速的推理服務。 AITemplate 的亮點包括：

• 高性能：在主要模型上接近roofline fp16 TensorCore（NVIDIA GPU）/MatrixCore（AMD GPU）性能，包括ResNet、MaskRCNN、BERT、VisionTransformer、Stable Diffusion 等。
• 統一、開放、靈活：用於NVIDIA GPU 或AMD GPU 的Seamless fp16 深度神經網絡模型。完全開放源程式碼，樂高式的易擴展高性能基元，支持新的模型。

安裝

硬件要求：

• NVIDIA ：AIT 僅在SM80+ GPU 上進行測試，並非所有內核都適用於舊的SM75/SM70 (T4/V100) GPU。
• AMD ：AIT 僅在CDNA2 (MI-210/250) GPU 上進行測試，舊的CDNA1 (MI-100) GPU 可能存在編譯器問題。

克隆程式碼

克隆程式碼時，請使用以下命令同時克隆子模塊：

git clone --recursive https://github.com/facebookincubator/AITemplate

Docker 鏡像

我們強烈建議將AITemplate 與Docker 一起使用，以避免意外使用錯誤版本的NVCC 或HIPCC。

• CUDA: ./docker/build.sh cuda
• ROCM: DOCKER_BUILDKIT=1 ./docker/build.sh rocm

這將構建一個帶有 ait:latest標籤的docker 鏡像