AITemplateMeta 開源的 Python 框架

AITemplate（AIT）是一個(gè) Python 框架，它將深度神經(jīng)網(wǎng)絡(luò)轉(zhuǎn)化為 CUDA（NVIDIA GPU）/ HIP（AMD GPU）C++ 代碼，以實(shí)現(xiàn)快速的推理服務(wù)。AITemplate 的亮點(diǎn)包括：

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

安裝

硬件要求：

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

克隆代碼

克隆代碼時(shí)，請使用以下命令同時(shí)克隆子模塊：

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

Docker 鏡像

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

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

這將構(gòu)建一個(gè)帶有 ait:latest標(biāo)簽的 docker 鏡像