MAX78000是具有超低功耗卷積神經(jīng)網(wǎng)絡(luò)加速器的人工智能微控制器，可以在芯片上有效地運行人工智能模型。用戶應(yīng)首先使用ADI公司在PyTorch上的開發(fā)流程開發(fā)神經(jīng)網(wǎng)絡(luò)模型。然后，MAX78000頻率合成器工具接受YAML格式的PyTorch檢查點和模型描述，自動生成C代碼，在MAX78000上編譯和執(zhí)行。模型開發(fā)階段使用的基本軟件組件之一是數(shù)據(jù)加載器，它負責(zé)特定于應(yīng)用程序的數(shù)據(jù)準備任務(wù)。本文檔介紹在準備適合MAX78000模型訓(xùn)練的特定應(yīng)用訓(xùn)練和驗證/測試集實體時，數(shù)據(jù)加載器實現(xiàn)的原則和設(shè)計注意事項。

介紹

在應(yīng)用程序開發(fā)周期中，第一步是準備和預(yù)處理可用數(shù)據(jù)以創(chuàng)建訓(xùn)練和驗證/測試數(shù)據(jù)集。除了通常的數(shù)據(jù)預(yù)處理外，在MAX78000上運行模型還需要考慮幾個硬件限制。

數(shù)據(jù)加載器的主要職責(zé)可以總結(jié)如下：

[可選]將原始資源的輸入和標簽數(shù)據(jù)下載到通過調(diào)用ADI公司的CNN培訓(xùn)工具（培訓(xùn)存儲庫/train.py提供的數(shù)據(jù)路徑中。

從指定的數(shù)據(jù)路徑（csv/二進制文件、帶或不帶層次結(jié)構(gòu)的文件夾/s 等）讀取原始輸入數(shù)據(jù)。

讀取提供的數(shù)據(jù)路徑中的原始標簽/注釋（csv/二進制文件/s、帶或不帶層次結(jié)構(gòu)的文件夾等）。

[可選]應(yīng)用數(shù)據(jù)預(yù)處理步驟，如增強、數(shù)據(jù)清理等。

對輸入數(shù)據(jù)和標簽應(yīng)用所需的數(shù)據(jù)類型和范圍轉(zhuǎn)換。

執(zhí)行訓(xùn)練和測試/驗證拆分。

提供數(shù)據(jù)加載器方法和與MAX78000模型訓(xùn)練工具兼容的定義字典。

[可選]將處理后的數(shù)據(jù)實體保留在磁盤上，以便將來訪問。

[可選]將上述步驟應(yīng)用于可從同一原始數(shù)據(jù)源生成的每個不同數(shù)據(jù)集變體。

提供兩個用于訓(xùn)練和測試數(shù)據(jù)的 PyTorch 數(shù)據(jù)集。

以下各節(jié)提供了有關(guān)創(chuàng)建高效數(shù)據(jù)加載器的說明，以滿足所需的功能并方便地集成訓(xùn)練工具。

圖 1 抽象地顯示了數(shù)據(jù)加載器實現(xiàn)的主流。以下各節(jié)介紹了詳細信息。

圖1.數(shù)據(jù)加載器模塊的主流。

自定義數(shù)據(jù)加載器實現(xiàn)的設(shè)計原則

數(shù)據(jù)加載器實現(xiàn)的主要職責(zé)之一是在將數(shù)據(jù)集實體饋送到 CNN 模型之前進行數(shù)據(jù)范圍調(diào)整和數(shù)據(jù)類型管理。圖 2 總結(jié)了這些操作，以下各節(jié)將詳細介紹這些操作。

圖2.數(shù)據(jù)范圍規(guī)范化和類型轉(zhuǎn)換。

預(yù)期數(shù)據(jù)范圍

對于訓(xùn)練，輸入數(shù)據(jù)應(yīng)在

.當(dāng)評估量化權(quán)重或在硬件上運行時，輸入數(shù)據(jù)應(yīng)位于本機MAX7800X范圍[-128， +127]。

如以下部分所述，數(shù)據(jù)加載器函數(shù)將數(shù)據(jù)路徑和一些參數(shù)作為輸入?yún)?shù)。參數(shù)字段包括兩個必填字段：act_mode_8bit 和 truncate_testset。當(dāng)設(shè)置為True時，第一個參數(shù)是指對于本地MAX7800X范圍，即范圍[-128， +127]，應(yīng)正確進行大小寫歸一化。設(shè)置為 False 時，規(guī)范化應(yīng)在訓(xùn)練范圍內(nèi)

如果可用數(shù)據(jù)在 [0 1] 范圍內(nèi)，例如，在 PIL 圖像中，數(shù)據(jù)加載器可以直接調(diào)用 ai8x.normalize（） 函數(shù)，使用提供的 args 參數(shù)將數(shù)據(jù)規(guī)范化為兩個支持的數(shù)據(jù)范圍：

class normalize:
"""
Normalize input to either [-128/128, +127/128] or [-128, +127]
"""
    def __init__(self, args):
        self.args = args

    def __call__(self, img):
        if self.args.act_mode_8bit:
            return img.sub(0.5).mul(256.).round().clamp(min=-128, max=127)
        return img.sub(0.5).mul(256.).round().clamp(min=-128, max=127).div(128.)

如果可用數(shù)據(jù)范圍為 [0 255]，則需要在調(diào)用 ai256x.normalize（） 函數(shù)之前將其除以 0 以使其達到 [1 8] 范圍。

注意：ai8x 模塊的設(shè)備設(shè)置方法ai8x.set_device也接受相關(guān)參數(shù)模擬：True 表示訓(xùn)練案例 （act_mode_8bit = True），F(xiàn)alse 表示量化模型的評估或在也初始化 act_mode_8bit = False 的硬件上運行。此方法由具有適當(dāng)參數(shù)管理的訓(xùn)練腳本使用，但如果在外部調(diào)用函數(shù)，則應(yīng)正確設(shè)置模擬參數(shù)。

在MAX7800X硬件上運行推理時，必須考慮本地數(shù)據(jù)格式，并且在推理過程中應(yīng)盡可能少地進行預(yù)處理。

數(shù)據(jù)類型

數(shù)據(jù)源可能具有不同范圍內(nèi)各種格式和值的原始數(shù)據(jù)文件。數(shù)據(jù)集類和數(shù)據(jù)加載器函數(shù)負責(zé)處理必要的轉(zhuǎn)換。

數(shù)據(jù)加載器函數(shù)應(yīng)返回數(shù)據(jù)類的訓(xùn)練和測試數(shù)據(jù)集元組。類型轉(zhuǎn)換和轉(zhuǎn)換通常在 __get_item__ 函數(shù)內(nèi)處理，該函數(shù)應(yīng)返回指定索引數(shù)據(jù)實體的數(shù)據(jù)元組和標簽。數(shù)據(jù)項的類型應(yīng)為：火炬。庫達]。形狀火炬的浮動張量。大小（數(shù)據(jù)集字典的相關(guān)條目“輸入”字段）。

標簽維度可能因問題類型或輸入數(shù)據(jù)形狀而異。每個標簽類型都應(yīng)強制轉(zhuǎn)換為 np.long，以便在訓(xùn)練腳本中正確計算訓(xùn)練損失。

在完成所有數(shù)據(jù)增強和預(yù)處理任務(wù)并將數(shù)據(jù)范圍規(guī)范化為 [0 1] 后，應(yīng)使用該ai8x_normalize進行適當(dāng)?shù)倪M一步規(guī)范化，然后可以使用 torchvision.transforms.ToTensor 執(zhí)行類型轉(zhuǎn)換。

注意：拿到火炬。庫達]。FloatTensor，numpy 數(shù)組必須事先轉(zhuǎn)換為 float32。

Torchvision軟件包包括各種預(yù)處理轉(zhuǎn)換，例如可以根據(jù)應(yīng)用程序需求使用的PIL圖像的隨機裁剪。數(shù)據(jù)類可以利用Torchvision包的復(fù)合轉(zhuǎn)換按順序應(yīng)用多個轉(zhuǎn)換，例如ToTensor轉(zhuǎn)換和ai8x_normalize每當(dāng)訪問數(shù)據(jù)條目時。

數(shù)據(jù)實體的存儲

通常，有兩種方法可以存儲數(shù)據(jù)集條目;整個數(shù)據(jù)集條目可以存儲在內(nèi)存中，也可以在使用 __getitem__ 方法訪問時從磁盤讀取。基本的決策因素是數(shù)據(jù)集的大小和每個實體的大小。當(dāng)數(shù)據(jù)集太大而無法放入內(nèi)存時（在初始化函數(shù)中處理預(yù)處理和增強任務(wù)后），所有數(shù)據(jù)集條目都可以保存到磁盤中，并在以后每次訪問時從磁盤單獨讀取。雖然將數(shù)據(jù)條目保留在內(nèi)存中可以加快數(shù)據(jù)訪問速度，但內(nèi)存限制可能會阻止在所有情況下使用基于內(nèi)存的方法。

注意：即使采用基于內(nèi)存的方法，也建議將預(yù)處理和增強的數(shù)據(jù)條目寫入磁盤，因為它們只執(zhí)行一次。然后，在每次生成數(shù)據(jù)類實例時，可以執(zhí)行將所有數(shù)據(jù)批量讀取到內(nèi)存中。

表 1 總結(jié)了同一數(shù)據(jù)源的兩個數(shù)據(jù)加載器實現(xiàn)選項的一些度量。從第一行可以看出，磁盤存儲方法可以處理更多圖像。內(nèi)存預(yù)算是限制已處理圖像數(shù)量的因素。這兩種方法的數(shù)據(jù)集生成時間都很長，因為第一種方法還處理預(yù)處理、擴充等步驟，然后將所有可用數(shù)據(jù)寫入磁盤。基于內(nèi)存的方法在以后生成數(shù)據(jù)集需要更長的時間，因為對象創(chuàng)建需要將大文件從磁盤批量讀取到內(nèi)存中。而在第二種方法中，每個數(shù)據(jù)集項都是獨立保存的，__getitem__方法創(chuàng)建數(shù)據(jù)集對象和實體檢索都花費很少的時間。第一種方法的內(nèi)存消耗很高，因為它將所有數(shù)據(jù)集實體保留在內(nèi)存中。在磁盤使用情況方面，第一種方法通常使用所有數(shù)據(jù)條目的單個文件，第二種方法為每個數(shù)據(jù)條目使用單獨的文件。這應(yīng)該會導(dǎo)致大致相似的磁盤預(yù)算。

注意：在表 1 中，由于處理圖像的數(shù)量減少，第一種方法的磁盤空間要小得多。磁盤方法的唯一缺點是它增加了訓(xùn)練時間，因為每個數(shù)據(jù)輸入讀取都是作為單獨的磁盤操作完成的。

自定義數(shù)據(jù)加載器實現(xiàn)的編程原則

數(shù)據(jù)加載器模塊將在 PyTorch 中實現(xiàn)，預(yù)計至少具有以下三個組件：

數(shù)據(jù)集類定義。示例：類 AISegment（數(shù)據(jù)集）

torch.utils.data.Dataset是自定義數(shù)據(jù)集實現(xiàn)類應(yīng)繼承的抽象類。有關(guān) PyTorch 中自定義數(shù)據(jù)加載器實現(xiàn)的教程，請參閱 [1]。

應(yīng)重寫__len__方法，以便 len（dataset） 返回數(shù)據(jù)集的大小。

還應(yīng)該實現(xiàn)__getitem__來支持索引，以便數(shù)據(jù)集[i]可用于獲取i千樣本。對于MAX78000應(yīng)用，該方法應(yīng)返回一個數(shù)據(jù)元組及其相應(yīng)的標簽。

__init__功能參數(shù)和內(nèi)容可根據(jù)應(yīng)用需求進行定制。前兩個參數(shù)通常是數(shù)據(jù)根路徑和類型（測試或訓(xùn)練），如MAX78000訓(xùn)練存儲庫數(shù)據(jù)集文件夾中的幾個數(shù)據(jù)加載器實現(xiàn)所示。但是，只要以下項目中介紹并作為外部通信點提供的數(shù)據(jù)加載器函數(shù)是固定的并執(zhí)行所需的操作，就可以更改這些參數(shù)的順序或命名。

數(shù)據(jù)加載器函數(shù)：不應(yīng)修改此函數(shù)的簽名。第一個輸入是指定數(shù)據(jù)目錄和程序參數(shù)的元組。兩個前導(dǎo)布爾輸入指定是否應(yīng)加載訓(xùn)練和/或測試數(shù)據(jù)。

程序參數(shù)有兩個與數(shù)據(jù)類實現(xiàn)相關(guān)的關(guān)鍵字段;act_mode_8bit和truncate_testset。第一個是指規(guī)范化類型（有關(guān)更多詳細信息，請參閱預(yù)期數(shù)據(jù)范圍部分），第二個用于將測試集截斷為單個元素集。

示例：def AISegment352_get_datasets（data， load_train=True， load_test=True）。

數(shù)據(jù)集字典包括可用的數(shù)據(jù)加載器函數(shù)。字典名稱和鍵值都不應(yīng)更改，僅應(yīng)根據(jù)自定義數(shù)據(jù)集實現(xiàn)調(diào)整值。同一數(shù)據(jù)源的每個變體都可以作為此字典中的單獨元素存在。

例：

“name”鍵的值使ADI公司的CNN訓(xùn)練工具（訓(xùn)練存儲庫/train.py）能夠在提供--dataset參數(shù)時查找數(shù)據(jù)集。因此，此字段的值在自定義數(shù)據(jù)集中應(yīng)該是唯一的。

“input”鍵的值是輸入數(shù)據(jù)的維度。第一個維度作為num_channels傳遞給模型，而其余維度作為維度傳遞。例如，“input”：（1， 28， 28） 以 num_channels=1 和 dimension=（28， 28） 的形式傳遞給模型。一維輸入使用單個“維度”，例如“input”：（2， 512） 以 num_channels=2 和 dimension=（512， ） 的形式傳遞給模型。

“output”鍵的值指定分類問題的可用類類型。此鍵的值也可以使用字符串文本定義。

示例：“輸出”：（“背景”、“縱向”）。

“weight”鍵的值指定每個數(shù)據(jù)實體的權(quán)重，并引用類標簽。這是一個可選字段，如果未提供，則全部為“1”。

可以通過使用可用樣本數(shù)為每個類提供成反比的權(quán)重來解決訓(xùn)練數(shù)據(jù)集中的類不平衡問題。因此，訓(xùn)練腳本更加關(guān)注頻率較低的樣本。

可選的回歸可以設(shè)置為 True 以自動選擇訓(xùn)練腳本的 --regression 命令行參數(shù)。

注意：當(dāng)類數(shù)為 1 時，訓(xùn)練腳本會自動設(shè)置回歸。 示例：“輸出”：（“id”），“回歸”：真。

示例數(shù)據(jù)加載器

MAX78000訓(xùn)練存儲庫數(shù)據(jù)集文件夾包括幾種不同的數(shù)據(jù)加載器實現(xiàn)，詳情請參見[2]。在本節(jié)中，將介紹一個定制的數(shù)據(jù)加載器來舉例說明所有提到的原則?？v向分割數(shù)據(jù)集用于此目的，更多詳細信息請參見 [3]。此數(shù)據(jù)集源包括 34，427 張分辨率為 600 × 800 的人類肖像圖像（紅色、綠色和藍色 （RGB） 顏色格式），以及相同數(shù)量的標簽圖像，具有相同大小的紅色、綠色、藍色和 Alpha （RGBA） 格式的相應(yīng)蒙版。

初始化

設(shè)計的數(shù)據(jù)加載器模塊的第一個組件是具有以下初始化函數(shù)的數(shù)據(jù)加載器類。將跳過有關(guān)生成數(shù)據(jù)集信息數(shù)據(jù)框的詳細信息。簡而言之，這些行包括一些路徑處理代碼，用于保留原始圖像路徑、原始遮罩文件路徑、裁剪 idx 和要保存到的數(shù)據(jù)集條目的泡菜文件路徑。除了這些路徑生成部分之外，初始化函數(shù)的主要功能是保留提供的參數(shù)并相應(yīng)地排列一些局部變量（例如，訓(xùn)練或測試數(shù)據(jù)集信息數(shù)據(jù)幀）并生成數(shù)據(jù)集實體。

對于第一次初始化調(diào)用，所有數(shù)據(jù)處理任務(wù)都使用 __gen_datasets__ 方法處理，并為每個數(shù)據(jù)集項生成 pickle 文件并存儲在磁盤上，以便在每次數(shù)據(jù)訪問時讀取。

class AISegment(Dataset):
…
    def __init__(self, root_dir, d_type, transform=None, im_size=[80, 80], fold_ratio=1):
        …
        self.d_type = d_type
        self.transform = transform
        self.img_ds_dim = im_size
        self.fold_ratio = fold_ratio

        # Generate and save dataset information file if not already available
        # Training and Test split is also performed here using the hash of file names (all three cropped images should fall into the same set)
        # Information data frames include raw data path, raw label path, crop idx, pickle file path, etc. for each data entity
        …

        # One of the created data frames is selected from: train_img_files_info & test_img_files_info
        if self.d_type == 'train':
            self.img_files_info = train_img_files_info

        elif self.d_type == 'test':
            self.img_files_info = test_img_files_info
        else:
            print('Unknown data type: %s' % self.d_type)
            return

   # Create and save pt files for each data entity (if not available before)
   self.__create_pt_files()
        self.is_truncated = False

    def __create_pt_files(self):
        if self.__check_pt_files_exist():
            return
        self.__makedir_exist_ok(self.processed_train_data_folder)
        self.__makedir_exist_ok(self.processed_test_data_folder)
        self.__gen_datasets()

數(shù)據(jù)增強

gen_datasets方法處理所有必需的預(yù)處理、擴充和預(yù)規(guī)范化步驟。實施的步驟如下：

從原始圖像裁剪三個正方形圖像（因為U-Net模型使用方形圖像）。

根據(jù)提供的數(shù)據(jù)集參數(shù)，裁剪圖像和遮罩圖像被下采樣為 80×80 或 352×352。

相應(yīng)的遮罩圖像將轉(zhuǎn)換為二進制“背景”或“肖像”標簽。

如果需要，將折疊圖像（352×352 圖像折疊為大小為 88×88×48 的圖像）。

圖像在保存之前按 256 縮放，因為復(fù)合變壓器期望圖像在 [0 1] 范圍內(nèi)，但原始圖像是 RGB 并且值在 [0 255] 范圍內(nèi)。

圖像裁剪原理如圖3所示。

圖 4 包括從同一原始圖像構(gòu)建的三個示例訓(xùn)練圖像。

圖3.示例數(shù)據(jù)加載器 – 數(shù)據(jù)增強：從原始圖像中裁剪出三個方形圖像。

圖4.從原始圖像（600 × 600）裁剪了三個大小為 800 × 600 的圖像（以及相應(yīng)的消光圖像）。

gen_datasets方法的實現(xiàn)方式如下：

def __normalize_image(self, image):
        return image / 256

def __gen_datasets(self):

        # For each entry in dataset information dataframe
        for _, row in tqdm(self.img_files_info.iterrows()):

            img_file = row['img_file_path']
            matting_file = row['lbl_file_path']
            pickle_file = row['pickle_file_path']
            img_crp_idx = row['crp_idx']

            img = Image.open(img_file)
            lbl_rgba = Image.open(matting_file)

            vertical_crop_area = AISegment.img_dim[0] - AISegment.img_crp_dim[0]
            step_size = vertical_crop_area / (AISegment.num_of_cropped_imgs - 1)

            # Determine top left coordinate of the crop area
            top_left_x = 0
            top_left_y = 0 + img_crp_idx * step_size
                
            # Determine bottom right coordinate of the crop area
            bottom_right_x = AISegment.img_crp_dim[0]
            bottom_right_y = top_left_y + AISegment.img_crp_dim[0]
                
            img_crp = img.crop((top_left_x, top_left_y, bottom_right_x, bottom_right_y))
            img_crp_lbl = lbl_rgba.crop((top_left_x, top_left_y, bottom_right_x, bottom_right_y))
    
            img_crp = img_crp.resize(self.img_ds_dim)
            img_crp = np.asarray(img_crp).astype(np.uint8)

            img_crp_lbl = img_crp_lbl.resize(self.img_ds_dim)
            img_crp_lbl = (np.asarray(img_crp_lbl)[:, :, 3] == 0).astype(np.uint8) 
                
            # Fold the data (ex: 352 x 352 x 3 folded into 88 x 88 x 48) if required and save to pt file
            if self.fold_ratio == 1:
                img_crp_folded = img_crp

            else:
                img_crp_folded = None
                for i in range(self.fold_ratio):
                    for j in range(self.fold_ratio):
                        if img_crp_folded is not None:
                            img_crp_folded = np.concatenate((img_crp_folded, img_crp[i::self.fold_ratio, j::self.fold_ratio, :]), axis=2)
                        else:
                            img_crp_folded = img_crp[i::self.fold_ratio, j::self.fold_ratio, :]
            
            pickle.dump((img_crp_folded, img_crp_lbl), open(pickle_file, 'wb'))

數(shù)據(jù)加載器方法和轉(zhuǎn)換器定義

數(shù)據(jù)加載器方法是第二個必需的組件定制數(shù)據(jù)加載器模塊。對于示例 AISegment 數(shù)據(jù)集，實現(xiàn)了兩個不同的數(shù)據(jù)加載器函數(shù)。第一個 （AISegment_get_datasets） 使用較小的 U-Net 網(wǎng)絡(luò)模型返回大小為 80x80 的圖像。后者 （AISegment352_get_datasets） 返回大小為 352×352 的圖像。下面是第二個的實現(xiàn)，它生成具有所需屬性的 AISegment 對象。復(fù)合變壓器也在此函數(shù)中定義。此外，如果需要截斷，則測試數(shù)據(jù)集將被截斷。

def AISegment352_get_datasets(data, load_train=True, load_test=True):
    """…"""
    (data_dir, args) = data

    if load_train:
        train_transform = transforms.Compose([
            transforms.ToTensor(),
            ai8x.normalize(args=args)
        ])
        
        train_dataset = AISegment(root_dir=data_dir, d_type='train',
                                  transform=train_transform,
                                  im_size=[352, 352])
    else:
        train_dataset = None

    if load_test:
        test_transform = transforms.Compose([
            transforms.ToTensor(),
            ai8x.normalize(args=args)
        ])

        test_dataset = AISegment(root_dir=data_dir, d_type='test',
                                 transform=test_transform,
                                 im_size=[352, 352])

        if args.truncate_testset:
            test_dataset.data = test_dataset.data[:1]
    else:
        test_dataset = None

    return train_dataset, test_dataset

數(shù)據(jù)集字典

數(shù)據(jù)集字典是自定義數(shù)據(jù)加載器模塊的第三個必需組件，該模塊包含可用的數(shù)據(jù)加載器功能。對于示例 AISegment 數(shù)據(jù)加載程序，由于有兩種數(shù)據(jù)集變體可以生成具有不同分辨率（80×80 或 352×352）的數(shù)據(jù)集實體，因此數(shù)據(jù)集字典有兩個實體，每個實體都包括輸入和輸出大小的正確定義以及數(shù)據(jù)加載器函數(shù)名稱。

使用圖像測試訓(xùn)練模型

在模型開發(fā)階段之后，可以使用測試數(shù)據(jù)集或任意樣本來測試模型。關(guān)鍵的一點是，在向模型提供輸入之前，必須在外部完成正確的轉(zhuǎn)換/s數(shù)據(jù)加載器實現(xiàn)。

例如，使用AISegment數(shù)據(jù)集訓(xùn)練的示例模型需要輸入形狀[48， 88， 88]，這是分辨率為352×352的折疊RGB圖像的通道優(yōu)先表示，具有MAX7800X所需的歸一化像素值。外部提供的測試圖像甚至可能沒有相同的顏色格式，但必須事先實現(xiàn)所需的轉(zhuǎn)換，因為模型是針對 RGB 圖像訓(xùn)練的。下面是具有 470×470 分辨率和 YbCr 顏色格式的縱向圖像的測試模型的示例代碼片段：

rgb_img = yuv_img.convert('RGB')
rgb_img_ds = rgb_img.resize([352, 352])

# Image to numpy array conversion:
rgb_img_ds = np.asarray(rgb_img_ds).astype(np.uint8)

# Fold image (352 x 352 x 3 folded into 88 x 88 x 48)
rgb_img_ds_folded = fold_image(rgb_img_ds, 4)

# Covert pixel values to range [0 1] and cast to float type (required for Torch)
rgb_img_ds_folded_scaled = (rgb_img_ds_folded / 256).astype(np.float32)

# Normalize for MAX78000
# Set act_mode_8bit=True as we will set model parameter simulate=True     
args = Args(act_mode_8bit=True)

transform = transforms.Compose([
            transforms.ToTensor(),
            ai8x.normalize(args=args)
])
rgb_img_ds_folded_scaled_normalized = transform(rgb_img_ds_folded_scaled)

# Add batch dimension
rgb_img_batch = rgb_img_ds_folded_scaled_normalized.unsqueeze(0)

# Load model
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
load_model_path = 'unet/qat_ai85unet_v7_352_4_best_q.pth.tar'

ai8x.set_device(device=85, simulate=True, round_avg=False)

model = mod.AI85Unet_v7_pt(num_classes=2, num_channels=3, dimensions=(88, 88), 
                        bias=True, fold_ratio=4)

checkpoint = torch.load(load_model_path, map_location=lambda storage, loc: storage)
ai8x.fuse_bn_layers(model)
model = apputils.load_lean_checkpoint(model, load_model_path, model_device=device)
ai8x.update_model(model)
model = model.to(device)

# Run model
with torch.no_grad():
    sample_img_rgb_batch = rgb_img_batch.to(device)
    model_out_rgb = model(sample_img_rgb_batch)

# Retrieve model output
out_vals = np.argmax(model_out_rgb[0, :, :, :].detach().cpu().numpy(), axis=0)
plt.imshow(out_vals, cmap='Greys')

圖5包括以YCbCr格式給出的示例外部測試數(shù)據(jù)項，相應(yīng)的RGB圖像以及執(zhí)行所有必需轉(zhuǎn)換后的模型輸出。首先，色彩空間需要轉(zhuǎn)換為RGB。然后，應(yīng)將圖像縮減像素采樣，使其具有 352 × 352 分辨率。下一個操作是折疊，需要轉(zhuǎn)換和規(guī)范化。

圖5.在 YCbCr 色彩空間、RGB 空間和模型輸出中采樣縱向圖像。

審核編輯：郭婷