隨著短視頻 APP、視頻會(huì)議平臺(tái)以及 VR/AR 等技術(shù)的發(fā)展，視頻與圖像已逐漸成為全球互聯(lián)網(wǎng)流量的主要組成部分。包含我們平時(shí)接觸到的這些視頻圖像，也有很多是被 AI 和計(jì)算機(jī)視覺（CV）算法處理并增強(qiáng)過的。然而，隨著社交媒體和視頻分享服務(wù)的快速增長，作為 AI 圖像算法基礎(chǔ)的視頻圖像處理部分，也早已成為計(jì)算流程中不可忽視的成本和瓶頸。

CV-CUDA 應(yīng)用場(chǎng)景

我們先帶大家簡單回顧一下圖像處理的一些常見的例子，以更好地理解 CV-CUDA（Computer Vision - Compute Unified Device Architecture）應(yīng)用場(chǎng)景。

（1）AI 算法圖像背景模糊化

圖 1. AI 背景模糊（CPU 前后處理方案）

圖像背景模糊化通常會(huì)被應(yīng)用于視頻會(huì)議，美圖修圖等場(chǎng)景。在這些場(chǎng)景中，我們通常希望 AI 算法可以把主體之外的背景部分模糊化，這樣可以保護(hù)用戶隱私，美化圖像等。圖像背景模糊化的流程大體可以分為 3 個(gè)過程：前處理，DNN 網(wǎng)絡(luò)以及后處理過程。前處理過程，通常包含了對(duì)圖像做 Resize、Padding、Image2Tensor 等操作；DNN 網(wǎng)絡(luò)可以是一些常見 segmentation network，比如 Unet 等；后處理過程，通常包括 Tensor2Mask、Crop、Resize、Denoise 等操作。

在傳統(tǒng)的圖像處理流程中，前處理和后處理部分通常都是使用 CPU 進(jìn)行操作，這導(dǎo)致整個(gè)圖像背景模糊化流程中，有 90% 的工作時(shí)間消耗在前后處理部分，因而成為了整個(gè)算法流水線的瓶頸。若能把前后處理妥善利用 GPU 加速，這將能大幅提升整體的計(jì)算性能。

圖 2. AI 背景模糊（GPU 前后處理方案）

當(dāng)我們把前后處理部分都放到 GPU 上后，我們就可以對(duì)整個(gè) pipeline 進(jìn)行端到端的 GPU 加速。經(jīng)過測(cè)試，在單個(gè) GPU 上，相比于傳統(tǒng)圖像處理方式，把整個(gè) pipeline 移植到 GPU 后，可以獲得 20 倍以上的吞吐率提升。這無疑會(huì)大大的節(jié)省計(jì)算成本。

（2）AI 算法圖像分類

圖 3. AI 圖像分類

圖像分類是最常見的 AI 圖像算法之一，通?？梢杂糜谖矬w識(shí)別，以圖搜圖等場(chǎng)景，幾乎是所有 AI 圖像算法的基礎(chǔ)。圖像分類的 pipeline 大體可以分為 2 個(gè)部分：前處理部分和 DNN 部分。其中前處理部分，在訓(xùn)練和推理過程中最常見的 4 種操作包括：圖片解碼、Resize、Padding、Normalize。DNN 部分已經(jīng)有了 GPU 的加速，而前處理部分通常都會(huì)使用 CPU 上的庫函數(shù)進(jìn)行處理。如果能夠把前處理部分也移植到 GPU 上，那么一方面可以釋放 CPU 資源，另一方面也可以進(jìn)一步提升 GPU 利用率，從而可以對(duì)整個(gè) pipeline 進(jìn)行加速。

當(dāng)前圖像處理主流方案

針對(duì)前后處理部分，我們總結(jié)一下目前已有的一些主流應(yīng)用方案：

（1）目前應(yīng)用最廣泛的圖像處理庫是 OpenCV

前后處理中的操作，往往是依賴于任務(wù)類型的，這導(dǎo)致了前后處理相關(guān)的操作種類繁多，數(shù)量龐雜。利用 OpenCV 確實(shí)可以實(shí)現(xiàn)大部分的圖像處理操作，不過，這些操作大部分是用 CPU 實(shí)現(xiàn)的，缺少對(duì)應(yīng)的 GPU 加速版本。而部分 CPU 操作雖然存在對(duì)應(yīng)的 GPU 版本，但是這些 GPU 版本的實(shí)現(xiàn)可能存在一些問題，包括：

a. 部分算子的 CPU 和 GPU 結(jié)果無法對(duì)齊。在某些場(chǎng)景下，我們的算法團(tuán)隊(duì)使用了 CPU 的算子進(jìn)行訓(xùn)練，而在推理階段負(fù)責(zé)部署的同學(xué)考慮到性能問題，決定使用 GPU 算子，然而在測(cè)試過程中，發(fā)現(xiàn) CPU 和 GPU 算子結(jié)果無法對(duì)齊，這可能會(huì)導(dǎo)致端到端的處理結(jié)果出現(xiàn)精度異常。

b. 部分算子 GPU 性能比 CPU 性能還弱。除了某些算法本身就不太適合 GPU 化之外，我們還發(fā)現(xiàn)部分 GPU 實(shí)現(xiàn)并不是很高效。比如，有些算子會(huì)在運(yùn)行過程中，臨時(shí)申請(qǐng)顯存空間，導(dǎo)致增加了額外的耗時(shí)。

c. 在真實(shí)的部署 pipeline 中，可能有些算子只有 CPU 版本，而有些算子既有 CPU 版本又有 GPU 版本。如果 pipeline 里既有 CPU 算子又有 GPU 算子，那么會(huì)存在 GPU 和 CPU 同步，以及 H2D，D2H 內(nèi)存拷貝問題，從而引入額外的耗時(shí)。

（2）使用 PyTorch 框架進(jìn)行模型訓(xùn)練引入的 torchvision 圖像處理庫

在算法開發(fā)階段，出于訓(xùn)練考慮，算法工程師可能會(huì)優(yōu)先使用 torchvision 來完成圖像處理操作。然而在模型部署階段，負(fù)責(zé)模型部署的工程師通常會(huì)使用 C++ 作為開發(fā)語言，沒有辦法使用 torchvision 完成圖像處理操作，那么可能會(huì)使用其他的 C++ 圖像處理庫，比如 OpenCV，然而這就涉及到，torchvision 和 OpenCV 結(jié)果對(duì)齊的問題。此外，與使用 OpenCV 的情況類似，torchvision 中的算子也有不支持 GPU 加速，或者 GPU 加速效果不佳的情況。

CV-CUDA 特點(diǎn)

如上所述，傳統(tǒng)的圖像預(yù)處理操作一般在 CPU 上進(jìn)行，一方面會(huì)占用大量的 CPU 資源，使得 CPU 和 GPU 的負(fù)載不均衡；另一方面由于基于 CPU 的圖像加速庫不支持 batch 操作，導(dǎo)致預(yù)處理的效率低下。為了解決當(dāng)前主流的圖像處理庫所存在的一些問題，NVIDIA 和字節(jié)跳動(dòng)的機(jī)器學(xué)習(xí)團(tuán)隊(duì)聯(lián)合開發(fā)了基于 GPU 的圖像處理加速庫 CV-CUDA，并擁有以下特點(diǎn)：

（1）Batch

支持 batch 操作，可以充分利用 GPU 高并發(fā)、高吞吐的并行加速特性，提升計(jì)算效率和吞吐率。

（2）Variable Shape

支持同一 batch 中圖片尺寸各不相同，保證了使用上的靈活性。此外，在對(duì)圖片進(jìn)行處理時(shí)，可以對(duì)每張圖片指定不同的參數(shù)。例如調(diào)用 RotateVarShape 算子時(shí)，可以對(duì) batch 中每張圖片指定不同的旋轉(zhuǎn)角度。

（3）C/C++/Python 接口

在部署機(jī)器學(xué)習(xí)算法時(shí)需要對(duì)齊訓(xùn)練和推理流程。一般來說，訓(xùn)練時(shí)利用 python 進(jìn)行快速驗(yàn)證，推理時(shí)利用 C++ 進(jìn)行高性能部署，然而一些圖像處理庫僅支持 python，這給部署帶來了極大的不便。如果在訓(xùn)練和推理采用不同的圖像處理庫，則又需要在推理端重新實(shí)現(xiàn)一遍邏輯，過程會(huì)非常繁瑣。

CV-CUDA 提供了 C、C++ 和 Python 接口，可以同時(shí)服務(wù)于訓(xùn)練和推理場(chǎng)景。從訓(xùn)練遷移到推理場(chǎng)景時(shí)，也可免去繁瑣的對(duì)齊流程，提高部署效率。

（4）獨(dú)立算子設(shè)計(jì)

CV-CUDA 作為基礎(chǔ)圖像處理庫，采用了獨(dú)立算子設(shè)計(jì)，不需要預(yù)先定義流水線。獨(dú)立算子的設(shè)計(jì)具有更高的靈活性，使調(diào)試變得更加的容易，而且可以使其與其他的圖像處理交互，或者將其集成在用戶自己的圖像處理上層框架中。

（5）結(jié)果對(duì)齊 OpenCV

不同的圖像處理庫由于對(duì)一些算子的實(shí)現(xiàn)方式不一致導(dǎo)致計(jì)算結(jié)果難以對(duì)齊。例如常見的 Resize 操作，OpenCV、OpenCV-gpu 以及 torchvision 的實(shí)現(xiàn)方式都不一樣，計(jì)算結(jié)果存在差異。因此如果在訓(xùn)練時(shí)用 OpenCV CPU 版本而推理時(shí)若要采用 GPU 版本或其他圖像處理庫，就會(huì)面臨結(jié)果存在誤差的問題。

在設(shè)計(jì)之初，我們考慮到當(dāng)前圖像處理庫中，很多用戶習(xí)慣使用 OpenCV 的 CPU 版本，因此在設(shè)計(jì)算子時(shí)，不管是函數(shù)參數(shù)還是圖像處理結(jié)果上，盡可能對(duì)齊 OpenCV CPU 版本的算子。當(dāng)用戶從 OpenCV 遷移到 CV-CUDA 時(shí)，只需做少許改動(dòng)便可使用，且圖片處理結(jié)果和 OpenCV 一致，不需要重新訓(xùn)練模型。

（6）易用性

CV-CUDA 提供了 Image、ImageBatchVarShape 等結(jié)構(gòu)體，方便用戶的使用。同時(shí)還提供了 Allocator 類，用戶可以自定義顯存分配策略（例如用戶可以設(shè)計(jì)顯存池分配策略來提高顯存分配速度），方便上層框架集成和管理資源。目前 CV-CUDA 提供了 PyTorch、OpenCV 和 Pillow 的數(shù)據(jù)轉(zhuǎn)化接口，方便用戶進(jìn)行算子替換和進(jìn)行不同圖像庫之間的混用。

（7）針對(duì)不同 GPU 架構(gòu)的性能高度優(yōu)化

CV-CUDA 可以支持 Volta、Turing、Ampere 等 GPU 架構(gòu)，并針對(duì)不同架構(gòu) GPU 的特點(diǎn)，在 CUDA kernel 層面進(jìn)行了性能上的高度優(yōu)化，可在云服務(wù)場(chǎng)景中規(guī)?；渴稹?/p>

算子數(shù)量及其性能

(1) CV-CUDA 目前提供了 50 多種算子

包括常用的 CvtColor、Resize、Crop、PadStack、Normalize 等。當(dāng)前算子能夠覆蓋大部分應(yīng)用場(chǎng)景，后續(xù)將會(huì)支持更多算子。

算子清單：

圖 4. CV-CUDA 支持算子清單

(2) 性能對(duì)比

在本文開頭的背景模糊算法里，采用 CV-CUDA 替代 OpenCV 和 TorchVision 的前后處理后，整個(gè)推理流程的吞吐率提升了 20 倍以上。下圖展示了在同一個(gè)計(jì)算節(jié)點(diǎn)上（2x Intel Xeon Platinum 8168 CPUs，1x NVIDIA A100 GPU），以 30fps 的幀率處理 1080p 視頻，采用不同的 CV 庫所能支持的最大的并行流數(shù)。測(cè)試采用了 4 個(gè)進(jìn)程，每個(gè)進(jìn)程 batchSize 為 64。

圖 5. AI 背景模糊（2x Intel Xeon Platinum 8168 CPUs , 1x NVIDIA A100 GPU）

其中涉及到的前處理操作有：

Resize (Downscale)、Padding、Convert Data Type、Normalize 及 Image to Tensor

涉及到的后處理操作有：

Tensor to Mask、Convert Data Type、Crop、Resize (Upscale)、Bilateral Filter (Denoise)、Gaussian Blur 及 Composite

對(duì)于單個(gè)算子的性能，我們也做了性能測(cè)試，下圖的測(cè)試場(chǎng)景選用的圖片大小為 480*360，CPU 選擇為 Intel(R) Core(TM) i9-7900X，BatchSize 大小為 1，進(jìn)程數(shù)為 1。

圖 6. CV-CUDA 性能對(duì)比（CPU: Intel(R) Core(TM) i9-7900X CPU @ 3.30 GHz）

性能優(yōu)化手段及例子

為了使 CV-CUDA 能夠更加高效的運(yùn)行在 GPU 上，我們采取了一系列的優(yōu)化手段。

（1）kernel 融合

采用了大量的 kernel 融合策略，減少了 kernel launch 和 global memory 的訪問時(shí)間。

（2）訪存優(yōu)化

采用了合并訪存，向量化讀寫，shared memory 等策略，提高了數(shù)據(jù)讀寫的效率。

（3）異步處理

CV-CUDA 中所有算子均采用異步處理的方式，可以減少同步帶來的等待耗時(shí)。

（4）高效計(jì)算

采用了 fast math、warp/block reduce、table lookup 等優(yōu)化手段，可以有效提升計(jì)算效率。

（5）預(yù)分配顯存

CV-CUDA 采用了預(yù)分配顯存策略，并且提供了 Allocator 類，幫助使用者自定義顯存分配策略或者可采取默認(rèn)的顯存分配策略。算子所需要的 buffer 和圖片顯存會(huì)在初始化階段分配好，而在執(zhí)行階段不會(huì)再進(jìn)行耗時(shí)的顯存分配操作。

框架和 API

（1）CV-CUDA 整體架構(gòu)

整個(gè) CV-CUDA 庫包含了以下幾個(gè)組成部分

a. CV-CUDA 核心模塊

核心模塊包含了 C/C++ 和 Python API、NVCV 模塊，Operator 算子模塊以及 CV-CUDA Tools。

b. CV-CUDA Interop 模塊

這個(gè)模塊包含了和其他圖像處理庫以及推理框架的交互接口，目前支持 OpenCV、Pytorch 和 Pillow，后續(xù)將陸續(xù)加入其他圖像處理庫的交互接口。

c. CV-CUDA Tools/Tests

包含一些單元測(cè)試模塊和工具函數(shù)

圖 7. CV-CUDA 整體架構(gòu)

（2）CV-CUDA 核心模塊

圖 8. CV-CUDA 核心模塊

核心模塊包含以下幾個(gè)部分

a. NVCV 核心支持庫

包含 Image 抽象類、Memory workspace 類、Allocator 和 Batch/VarShapeBatch 類

b. Operators CV 算子

包含了各種獨(dú)立算子（resize、filter 等等）

c. CV-CUDA Tools

包含了用于開發(fā)算子所需的各種功能函數(shù)

（3）API

圖 9. ImageBatchVarShape

圖 10. Resize

字節(jié)跳動(dòng)機(jī)器學(xué)習(xí)團(tuán)隊(duì)?wèi)?yīng)用案例

最初，CV-CUDA 項(xiàng)目的啟動(dòng)，主要目的是為了解決內(nèi)部模型訓(xùn)練和推理過程中預(yù)處理瓶頸問題。以模型訓(xùn)練為例，正如大家所熟知的，傳統(tǒng)的數(shù)據(jù)讀取和預(yù)處理過程都是由 CPU 完成的。當(dāng) CPU 多核算力足夠時(shí)，能為模型訓(xùn)練提供穩(wěn)定高效的數(shù)據(jù)，但當(dāng) CPU 算力不足時(shí)，數(shù)據(jù)準(zhǔn)備過程即會(huì)成為瓶頸，從而導(dǎo)致 GPU 利用率低下，造成 GPU 資源的浪費(fèi)。一直以來，工業(yè)界有不少模型的訓(xùn)練推理過程，都存在因?yàn)?CPU 資源競(jìng)爭而影響著性能的問題。這個(gè)問題在圖像類模型，尤其是針對(duì)圖像有比較復(fù)雜的預(yù)處理的場(chǎng)景中尤為明顯。而隨著 GPU 算力的不斷提升，CPU 資源競(jìng)爭將會(huì)越來越凸顯出來。為了解決這個(gè)問題，非常樸素的一個(gè)想法是，將造成 CPU 算力不足的預(yù)處理邏輯搬到 GPU 上來進(jìn)行，利用 GPU 龐大的算力，來高效的產(chǎn)生模型所需的數(shù)據(jù)。

但我們對(duì)這個(gè) GPU 預(yù)處理算子庫顯然是有要求的。首先是性能需要足夠好，如果預(yù)處理部分 GPU kernel 的實(shí)現(xiàn)跑起來非常慢，可能反而得不償失；其次預(yù)處理搬到 GPU 上做，對(duì)模型迭代本身的影響要小，我們的預(yù)期是預(yù)處理 kernel 是見縫插針?biāo)频?，在模型迭代過程中僅用一點(diǎn)點(diǎn) GPU 閑散資源即可；最后，預(yù)處理 kernel 需要支持定制化，字節(jié)跳動(dòng)機(jī)器學(xué)習(xí)團(tuán)隊(duì)內(nèi)部訓(xùn)練的模型多，需要的預(yù)處理邏輯也多種多樣，因此有許多定制的預(yù)處理邏輯需求。而 CV-CUDA 就很好的滿足了以上幾點(diǎn)的要求：每個(gè) op 支持 batch 數(shù)據(jù)輸入的設(shè)計(jì)，減少了頻繁 Kernel launch 的開銷，保證了高效的運(yùn)行性能；每個(gè) op 都支持 stream 對(duì)象和顯存指針的傳入，使得我們能更加靈活的配置相應(yīng)的 GPU 資源；每個(gè) op 設(shè)計(jì)開發(fā)時(shí)，既兼顧了通用性，也能按需提供定制化接口，能夠覆蓋內(nèi)部的圖片類預(yù)處理需求。

CV-CUDA 目前在抖音集團(tuán)內(nèi)部的多個(gè)線上線下場(chǎng)景得到了應(yīng)用，比如搜索多模態(tài)，圖片分類等，下面分模型訓(xùn)練和推理分別介紹一個(gè)典型應(yīng)用。（模型相關(guān)的細(xì)節(jié)這里就不多做介紹了，主要展示下 CV-CUDA 的應(yīng)用場(chǎng)景和獲得的收益。）

（1）模型訓(xùn)練

在字節(jié)跳動(dòng)機(jī)器學(xué)習(xí)團(tuán)隊(duì)內(nèi)部，CV-CUDA 在預(yù)處理邏輯復(fù)雜的訓(xùn)練任務(wù)上加速效果明顯。在模型訓(xùn)練過程中，數(shù)據(jù)準(zhǔn)備和模型迭代是個(gè)異步的過程，數(shù)據(jù)準(zhǔn)備只要不比模型迭代來的慢，其性能就是滿足需求的。但隨著新一代強(qiáng)力 GPU 的投入使用，模型本身的計(jì)算時(shí)間被進(jìn)一步縮短，因而對(duì)數(shù)據(jù)準(zhǔn)備的性能要求越來越高。此時(shí)預(yù)處理邏輯一復(fù)雜，CPU 競(jìng)爭就會(huì)變得激烈，這個(gè)時(shí)候 CV-CUDA 就能發(fā)揮很大的作用。

圖 11. 模型訓(xùn)練流程示意圖

字節(jié)跳動(dòng)機(jī)器學(xué)習(xí)團(tuán)隊(duì)某個(gè)視頻相關(guān)的多模態(tài)任務(wù)，CV-CUDA 應(yīng)用后就獲得了不小的性能收益。該任務(wù)預(yù)處理部分邏輯很復(fù)雜，既有多幀視頻的解碼，也有很多的數(shù)據(jù)增強(qiáng)，導(dǎo)致在 A100 上多卡訓(xùn)練時(shí) CPU 資源競(jìng)爭非常明顯，因而雖然已經(jīng)充分利用了 CPU（Intel Xeon Platinum 8336C）的多核性能，但數(shù)據(jù)預(yù)處理的速度仍然滿足不了模型計(jì)算的需求。而使用了 CV-CUDA，預(yù)處理邏輯被全部遷移到了 GPU，相應(yīng)的訓(xùn)練瓶頸也由數(shù)據(jù)預(yù)處理轉(zhuǎn)移到了模型計(jì)算本身，從而在模型訓(xùn)練的整體性能上我們獲得了近 90% 的收益。另一個(gè)獲得收益的 OCR 任務(wù)也十分相似，其預(yù)處理鏈路有十幾個(gè)算子，使用 CV-CUDA 后我們?cè)?V100 上訓(xùn)練即獲得了 80% 的加速。

圖 12. CV-CUDA 訓(xùn)練性能收益

需要注意的是，前后處理耗時(shí)占比不同的應(yīng)用在使用 CV-CUDA 后帶來的收益提升會(huì)有所差異。例如前后處理耗時(shí)占比為 50% 的應(yīng)用，使用 CV-CUDA 后，端到端吞吐率提升的理論上限為原始的 2 倍（理論上限是指預(yù)處理耗時(shí)為 0，端到端中模型的耗時(shí)占比為 100%，實(shí)際上前后處理加速后，依然會(huì)有耗時(shí)，因此理論上限僅用于分析理想情況能達(dá)到的最大加速比）；而前后處理占比 90% 以上的應(yīng)用，使用 CV-CUDA 后的吞吐率提升理論上限為原始的 10 倍以上。此外，CPU 性能的差異也會(huì)影響到最終的加速效果。

（2） 模型推理

模型推理和訓(xùn)練略有不同：因?yàn)槟Ｐ陀?xùn)練時(shí)數(shù)據(jù)讀取和模型更新是異步的，所以基本上提升模型訓(xùn)練性能可以與提高 GPU 利用率劃等號(hào)。但模型推理過程，數(shù)據(jù)讀取和模型計(jì)算是一個(gè)串行的過程，加速數(shù)據(jù)準(zhǔn)備過程本身即會(huì)對(duì)整體性能造成直接影響。再加上模型推理時(shí)只涉及到模型前向的計(jì)算，計(jì)算量一般都較小，GPU 通常有足夠的剩余資源來進(jìn)行其他的工作。因而模型推理階段，把預(yù)處理邏輯搬到 GPU 上來加速性能，既合理而又有可行性。

圖 13. 模型推理流程示意圖

字節(jié)跳動(dòng)機(jī)器學(xué)習(xí)團(tuán)隊(duì)的某個(gè)搜索多模態(tài)任務(wù)，在使用 CV-CUDA 之后，整體的上線吞吐相比于用 CPU 做預(yù)處理時(shí)有了 2 倍多的提升，值得一提的是，這里的基線本身即是經(jīng)過 CPU（Intel Xeon Platinum 8336C）多核等方式優(yōu)化過后的結(jié)果，雖然涉及到的預(yù)處理邏輯較簡單，但使用 CV-CUDA 之后加速效果依然非常明顯。

附錄

（1）開源計(jì)劃

（2）代碼示例

以圖片分類為例，以下代碼展示了如何利用 CV-CUDA 對(duì)圖片進(jìn)行預(yù)處理，以及如何和 Pytorch 進(jìn)行交互。

流程圖：

Python 代碼：

import torch
from torchvision import models
from torchvision.io.image import read_file, decode_jpeg
import numpy as np


# Import CV-CUDA module
import nvcv


"""
Image Classification python sample


The image classification sample uses Resnet50 based model trained on Imagenet
The sample app pipeline includes preprocessing, inference and post process stages
which takes as input a batch of images and returns the TopN classification results
of each image.


This sample gives an overview of the interoperability of pytorch with CVCUDA
tensors and operators
"""


# Set the image and labels file
filename = "./assets/tabby_tiger_cat.jpg"
labelsfile = "./models/imagenet-classes.txt"


# Read the input imagea file
data = read_file(filename)  # raw data is on CPU


# NvJpeg can be used to decode the image
# to the necessary color format on the device
inputImage = decode_jpeg(data, device="cuda")  # decoded image in on GPU
imageWidth = inputImage.shape[2]
imageHeight = inputImage.shape[1]


# decode_jpeg is currently limited to batchSize 1
# and Planar format (CHW)


# A torch tensor can be wrapped into a CV-CUDA Object using the "as_tensor"
# function in the specified layout. The datatype and dimensions are derived
# directly from the torch tensor.
nvcvInputTensor = nvcv.as_tensor(inputImage, "CHW")


# The input image is now ready to be used


# The Reformat operator can be used to convert CHW format to NHWC
# for the rest of the preprocessing operations


nvcvInterleavedTensor = nvcvInputTensor.reformat("NHWC")


"""
Preprocessing includes the following sequence of operations.
Resize -> DataType Convert(U8->F32) -> Normalize(Apply mean and std deviation)
-> Interleaved to Planar
"""


# Model settings
layerHeight = 224
layerWidth = 224
batchSize = 1


# Resize
# Resize to the input network dimensions
nvcvResizeTensor = nvcvInterleavedTensor.resize(
    (batchSize, layerWidth, layerHeight, 3), nvcv.Interp.CUBIC
)


# Convert to the data type and range of values needed by the input layer
# i.e uint8->float. A Scale is applied to normalize the values in the range 0-1
nvcvConvertTensor = nvcvResizeTensor.convertto(np.float32, scale=1 / 255)


"""
The input to the network needs to be normalized based on the mean and
std deviation value to standardize the input data.
"""


# Create a torch tensor to store the mean and standard deviation values for R,G,B
scale = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
scaleTensor = torch.Tensor(scale)
stdTensor = torch.Tensor(std)


# Reshape the the number of channels. The R,G,B values scale and offset will be
# applied to every color plane respectively across the batch
scaleTensor = torch.reshape(scaleTensor, (1, 1, 1, 3)).cuda()
stdTensor = torch.reshape(stdTensor, (1, 1, 1, 3)).cuda()


# Wrap the torch tensor in a CV-CUDA Tensor
nvcvScaleTensor = nvcv.as_tensor(scaleTensor, "NHWC")
nvcvBaseTensor = nvcv.as_tensor(stdTensor, "NHWC")


# Apply the normalize operator and indicate the scale values are std deviation
# i.e scale = 1/stddev
nvcvNormTensor = nvcvConvertTensor.normalize(
    nvcvBaseTensor, nvcvScaleTensor, nvcv.NormalizeFlags.SCALE_IS_STDDEV
)


# The final stage in the preprocess pipeline includes converting the RGB buffer
# into a planar buffer
nvcvPreprocessedTensor = nvcvNormTensor.reformat("NCHW")


# Inference uses pytorch to run a resnet50 model on the preprocessed input and outputs
# the classification scores for 1000 classes
# Load Resnet model pretrained on Imagenet
resnet50 = models.resnet50(pretrained=True)
resnet50.to("cuda")
resnet50.eval()


# Run inference on the preprocessed input
torchPreprocessedTensor = torch.as_tensor(nvcvPreprocessedTensor.cuda(), device="cuda")
inferOutput = resnet50(torchPreprocessedTensor)


"""
Postprocessing function normalizes the classification score from the network and sorts
the scores to get the TopN classification scores.
"""
# top results to print out
topN = 5


# Read and parse the classes
with open(labelsfile, "r") as f:
    classes = [line.strip() for line in f.readlines()]


# Apply softmax to Normalize scores between 0-1
scores = torch.nn.functional.softmax(inferOutput, dim=1)[0]


# Sort output scores in descending order
_, indices = torch.sort(inferOutput, descending=True)


# Display Top N Results
[
    print("Class : ", classes[idx], " Score : ", scores[idx].item())
    for idx in indices[0][:topN]
]

關(guān)于作者

張毅 | NVIDIA GPU 計(jì)算專家團(tuán)隊(duì)的工程師

主要從事計(jì)算機(jī)視覺領(lǐng)域的算法加速工作，同時(shí)也是 CV-CUDA 項(xiàng)目的發(fā)起人，負(fù)責(zé) CV-CUDA 項(xiàng)目早期的框架設(shè)計(jì)和算子開發(fā)工作。此外，他還是 PyTorch 轉(zhuǎn) TensorRT 編譯工具 Torch-TensorRT 的核心開發(fā)人員。在加入 NVIDIA 之前，于騰訊優(yōu)圖從事計(jì)算機(jī)視覺算法研發(fā)工作。分別于浙江大學(xué)和上海交通大學(xué)獲得學(xué)士和碩士學(xué)位。

郭若乾 | NVIDIA GPU 計(jì)算專家團(tuán)隊(duì)的工程師

于浙江大學(xué)取得碩士學(xué)位后，他在 2020 年加入 NVIDIA，主要從事計(jì)算機(jī)視覺領(lǐng)域的算法加速工作，以及 PyTorch 模型轉(zhuǎn) TensorRT 編譯工具 Torch-Tensor 的開發(fā)工作。作為 CV-CUDA 項(xiàng)目的主要開發(fā)人員，負(fù)責(zé)了 CV-CUDA 項(xiàng)目背景調(diào)研以及早期算子的開發(fā)優(yōu)化工作。

盛一耀 | 字節(jié)跳動(dòng)機(jī)器學(xué)習(xí)系統(tǒng)團(tuán)隊(duì)工程師

他是 CV-CUDA 項(xiàng)目主要的開發(fā)者之一。他于美國卡耐基梅隆大學(xué)獲得碩士學(xué)位后，2019 年加入字節(jié)跳動(dòng)，主要從事機(jī)器學(xué)習(xí)系統(tǒng)相關(guān)工作，對(duì)模型訓(xùn)練和推理性能優(yōu)化有較深刻的理解。在 CV-CUDA 項(xiàng)目中，他在把握并提供字節(jié)跳動(dòng)內(nèi)部視覺相關(guān)算子需求的同時(shí)，也參與開發(fā)了許多 GPU 視覺類算子。