首頁>技術>

當我在使用深度學習進行影象語義分割並想使用PyTorch在DeepLabv3[1]上執行一些實驗時,我找不到任何線上教程。並且torchvision不僅沒有提供分割資料集,而且也沒有關於DeepLabv3類內部結構的詳細解釋。然而,我是透過自己的研究進行了現有模型的遷移學習,我想分享這個過程,這樣可能會對你們有幫助。

在本文中,我將介紹如何使用預先訓練的語義分割DeepLabv3模型,透過使用遷移學習在PyTorch中進行道路裂縫檢測。同樣的過程也可以應用於調整自定義資料集的網路。

介紹

讓我們首先簡要介紹影象分割。 分割任務的主要目標是輸出畫素級輸出蒙版,其中將屬於某些類別的區域分配給相同的不同畫素值。 如果透過為每個類別分配不同的顏色來對這些細分蒙版進行顏色編碼以使其視覺化,那麼你就會得到一個類似於兒童塗色書中的影象。下面顯示了一個示例:

分割在計算機視覺和影象處理領域已經存在很長時間了。 其中一些技術是簡單的閾值化,基於聚類的方法,例如k均值聚類分割,區域增長方法等。[3]

隨著深度學習的最新進展以及卷積神經網路在影象相關任務中比傳統方法的成功,這些技術也已應用於影象分割任務。

遷移學習

當有限的資料可用時,深度學習模型往往會遇到困難。 對於大多數實際應用,即使不是不可能,也很難訪問大量資料集。 標註既繁瑣又費時。 即使您打算將其外包,您仍然必須花錢。 已經做出努力以能夠從有限的資料訓練模型。 這些技術中的一種稱為轉移學習。

遷移學習涉及使用針對源域和任務進行預訓練的網路(希望您可以在其中訪問大型資料集),並將其用於您的預期/目標域和任務(與原始任務和域類似) )[4]。 下圖可以從概念上表示它。

我們根據自己的要求更改目標細分子網路,然後訓練部分網路或整個網路。選擇的學習率低於正常訓練的學習率。 這是因為網路已經對源任務具有良好的權重。 我們不想太快地改變權重。有時也可以凍結初始層,因為有人認為這些層提取了一般特徵,可以潛在地使用而無需任何更改。

使用CrackForest資料集進行裂縫檢測

在本教程中,我將使用CrackForest [5] [6]資料集透過分段進行道路裂縫檢測。 它由具有裂縫作為缺陷的城市道路表面影象組成。 影象包含混淆區域,例如陰影,溢油和水漬。 這些影象是使用普通的iPhone5相機拍攝的。 資料集包含118張影象,並具有對應的裂紋畫素級別蒙版,所有蒙版的大小均為320×480。 額外的混雜因素以及可用於訓練的有限數量的樣本使CrackForest成為具有挑戰性的資料集[7]。

PyTorch的資料集

讓我們首先為模型構造一個數據集類,該資料集類將用於獲取訓練樣本。 為了進行分割,我們將一個地面真相掩碼影象作為標籤,而不是一個可以熱編碼的單值數字標籤。 蒙版具有可用的畫素級註釋,如圖3所示。因此,用於輸入和標籤的訓練張量將是四維的。 對於PyTorch,它們是:batch_size x通道x高x寬。

我們現在將定義細分資料集類。 類定義如下。

"""Author: Manpreet Singh MinhasContact: msminhas at uwaterloo ca"""from pathlib import Pathfrom typing import Any, Callable, Optionalimport numpy as npfrom PIL import Imagefrom torchvision.datasets.vision import VisionDatasetclass SegmentationDataset(VisionDataset):    """A PyTorch dataset for image segmentation task.    The dataset is compatible with torchvision transforms.    The transforms passed would be applied to both the Images and Masks.    """    def __init__(self,                 root: str,                 image_folder: str,                 mask_folder: str,                 transforms: Optional[Callable] = None,                 seed: int = None,                 fraction: float = None,                 subset: str = None,                 image_color_mode: str = "rgb",                 mask_color_mode: str = "grayscale") -> None:        """        Args:            root (str): Root directory path.            image_folder (str): Name of the folder that contains the images in the root directory.            mask_folder (str): Name of the folder that contains the masks in the root directory.            transforms (Optional[Callable], optional): A function/transform that takes in            a sample and returns a transformed version.            E.g, ``transforms.ToTensor`` for images. Defaults to None.            seed (int, optional): Specify a seed for the train and test split for reproducible results. Defaults to None.            fraction (float, optional): A float value from 0 to 1 which specifies the validation split fraction. Defaults to None.            subset (str, optional): 'Train' or 'Test' to select the appropriate set. Defaults to None.            image_color_mode (str, optional): 'rgb' or 'grayscale'. Defaults to 'rgb'.            mask_color_mode (str, optional): 'rgb' or 'grayscale'. Defaults to 'grayscale'.        Raises:            OSError: If image folder doesn't exist in root.            OSError: If mask folder doesn't exist in root.            ValueError: If subset is not either 'Train' or 'Test'            ValueError: If image_color_mode and mask_color_mode are either 'rgb' or 'grayscale'        """        super().__init__(root, transforms)        image_folder_path = Path(self.root) / image_folder        mask_folder_path = Path(self.root) / mask_folder        if not image_folder_path.exists():            raise OSError(f"{image_folder_path} does not exist.")        if not mask_folder_path.exists():            raise OSError(f"{mask_folder_path} does not exist.")        if image_color_mode not in ["rgb", "grayscale"]:            raise ValueError(                f"{image_color_mode} is an invalid choice. Please enter from rgb grayscale."            )        if mask_color_mode not in ["rgb", "grayscale"]:            raise ValueError(                f"{mask_color_mode} is an invalid choice. Please enter from rgb grayscale."            )        self.image_color_mode = image_color_mode        self.mask_color_mode = mask_color_mode        if not fraction:            self.image_names = sorted(image_folder_path.glob("*"))            self.mask_names = sorted(mask_folder_path.glob("*"))        else:            if subset not in ["Train", "Test"]:                raise (ValueError(                    f"{subset} is not a valid input. Acceptable values are Train and Test."                ))            self.fraction = fraction            self.image_list = np.array(sorted(image_folder_path.glob("*")))            self.mask_list = np.array(sorted(mask_folder_path.glob("*")))            if seed:                np.random.seed(seed)                indices = np.arange(len(self.image_list))                np.random.shuffle(indices)                self.image_list = self.image_list[indices]                self.mask_list = self.mask_list[indices]            if subset == "Train":                self.image_names = self.image_list[:int(                    np.ceil(len(self.image_list) * (1 - self.fraction)))]                self.mask_names = self.mask_list[:int(                    np.ceil(len(self.mask_list) * (1 - self.fraction)))]            else:                self.image_names = self.image_list[                    int(np.ceil(len(self.image_list) * (1 - self.fraction))):]                self.mask_names = self.mask_list[                    int(np.ceil(len(self.mask_list) * (1 - self.fraction))):]    def __len__(self) -> int:        return len(self.image_names)    def __getitem__(self, index: int) -> Any:        image_path = self.image_names[index]        mask_path = self.mask_names[index]        with open(image_path, "rb") as image_file, open(mask_path,                                                        "rb") as mask_file:            image = Image.open(image_file)            if self.image_color_mode == "rgb":                image = image.convert("RGB")            elif self.image_color_mode == "grayscale":                image = image.convert("L")            mask = Image.open(mask_file)            if self.mask_color_mode == "rgb":                mask = mask.convert("RGB")            elif self.mask_color_mode == "grayscale":                mask = mask.convert("L")            sample = {"image": image, "mask": mask}            if self.transforms:                sample["image"] = self.transforms(sample["image"])                sample["mask"] = self.transforms(sample["mask"])            return sample

我們使用torchvision中的VisionDataset類作為Segmentation資料集的基類。 以下三種方法需要過載。

init:此方法是資料集物件將初始化的位置。 通常,您需要構建影象檔案路徑和相應的標籤,它們是用於分割的遮罩檔案路徑。 然後,在lengetitem方法中使用這些路徑。 getitem:每當您使用object [index]訪問任何元素時,都會呼叫此方法。 因此,我們需要在此處編寫影象和蒙版載入邏輯。 因此,實質上,您可以使用此方法中的資料集物件從資料集中獲得一個訓練樣本。 len:每當使用len(obj)時,都會呼叫此方法。 此方法僅返回目錄中訓練樣本的數量。

為PyTorch建立自定義資料集時,請記住使用PIL庫。 這使您可以直接使用Torchvision轉換,而不必定義自己的轉換。

在此類的第一個版本中,我使用OpenCV來載入影象! 該庫不僅非常繁重,而且與Torchvision轉換不相容。 我必須編寫自己的自定義轉換並自己處理尺寸更改。

我添加了其他功能,使您可以將資料集保留在一個目錄中,而不是將Train和Val拆分到單獨的資料夾中,因為我使用的許多資料集都不採用這種格式,並且我不想重組我的資料集 資料夾結構每次。

現在我們已經定義了資料集類,下一步是從此建立一個PyTorch資料載入器。 資料載入器使您可以使用多執行緒處理來建立一批資料樣本和標籤。 這使得資料載入過程更加快捷和高效。 為此,可以使用torch.utils.data下可用的DataLoader類。 建立過程本身很簡單。 透過將資料集物件傳遞給它來建立一個DataLoader物件。 支援的引數如下所示。

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,           batch_sampler=None, num_workers=0, collate_fn=None,           pin_memory=False, drop_last=False, timeout=0,           worker_init_fn=None, *, prefetch_factor=2,           persistent_workers=False)

下面解釋了幾個有用的引數:

資料集(Dataset):要從中載入資料的資料集。

batch_size(整數,可選):每個批次要載入多少樣本(預設值:1)

shuffle(布林型,可選):設定為True可使資料在每個時期都重新隨機播放。 (預設值:False)

num_workers(int,可選):要用於資料載入的子程序數。 0表示將在主程序中載入資料。 (預設值:0)提示:您可以將此值設定為等於系統處理器中的核心數,以作為最佳值。 設定較高的值可能會導致效能下降。

此外,我編寫了兩個幫助程式函式,這些函式可以根據您的資料目錄結構為您提供資料載入器,並且可以在datahandler.py檔案中使用它們。

getdataloadersep_folder:從兩個單獨的Train和Test資料夾中建立Train和Test資料載入器。 目錄結構應如下所示。

data_dir --Train ------Image ---------Image1 ---------ImageN ------Mask ---------Mask1 ---------MaskN --Train ------Image ---------Image1 ---------ImageN ------Mask ---------Mask1 ---------MaskN

getdataloadersingle_folder:從單個資料夾建立。 結構應如下

--data_dir------Image---------Image1---------ImageN------Mask---------Mask1---------MaskN

接下來,我們討論本教程的關鍵所在,即如何根據我們的資料需求載入預訓練的模型並更改分割頭。

DeepLabv3模型

Torchvision有可用的預訓練模型,我們將使用其中一種模型。 我編寫了以下函式,該函式為您提供了具有自定義數量的輸出通道的模型。 如果您有多個班級,則可以更改此值。

""" DeepLabv3 Model download and change the head for your prediction"""from torchvision.models.segmentation.deeplabv3 import DeepLabHeadfrom torchvision import modelsdef createDeepLabv3(outputchannels=1):    """DeepLabv3 class with custom head    Args:        outputchannels (int, optional): The number of output channels        in your dataset masks. Defaults to 1.    Returns:        model: Returns the DeepLabv3 model with the ResNet101 backbone.    """    model = models.segmentation.deeplabv3_resnet101(pretrained=True,                                                    progress=True)    model.classifier = DeepLabHead(2048, outputchannels)    # Set the model in training mode    model.train()    return model

首先,我們使用models.segmentation.deeplabv3_resnet101方法獲得預訓練模型,該方法將預訓練模型下載到我們的系統快取中。注意resnet101是從此特定方法獲得的deeplabv3模型的基礎模型。這決定了傳遞到分類器的特徵向量的長度。

第二步是修改分割頭即分類器的主要步驟。該分類器是網路的一部分,負責建立最終的細分輸出。透過用具有新數量的輸出通道的新DeepLabHead替換模型的分類器模組來完成更改。 resnet101主幹的特徵向量大小為2048。如果您決定使用另一個主幹,請相應地更改此值。

最後,我們將模型設定為訓練模式。此步驟是可選的,因為您也可以在訓練邏輯中執行此操作。

下一步是訓練模型。

我定義了以下訓練模型的train_model函式。 它將訓練和驗證損失以及指標(如果指定)值儲存到CSV日誌檔案中,以便於訪問。 訓練程式碼程式碼如下。 後面有充分的文件來解釋發生了什麼。

def train_model(model, criterion, dataloaders, optimizer, metrics, bpath,                num_epochs):    since = time.time()    best_model_wts = copy.deepcopy(model.state_dict())    best_loss = 1e10    # Use gpu if available    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")    model.to(device)    # Initialize the log file for training and testing loss and metrics    fieldnames = ['epoch', 'Train_loss', 'Test_loss'] + \        [f'Train_{m}' for m in metrics.keys()] + \        [f'Test_{m}' for m in metrics.keys()]    with open(os.path.join(bpath, 'log.csv'), 'w', newline='') as csvfile:        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)        writer.writeheader()    for epoch in range(1, num_epochs + 1):        print('Epoch {}/{}'.format(epoch, num_epochs))        print('-' * 10)        # Each epoch has a training and validation phase        # Initialize batch summary        batchsummary = {a: [0] for a in fieldnames}        for phase in ['Train', 'Test']:            if phase == 'Train':                model.train()  # Set model to training mode            else:                model.eval()  # Set model to evaluate mode            # Iterate over data.            for sample in tqdm(iter(dataloaders[phase])):                inputs = sample['image'].to(device)                masks = sample['mask'].to(device)                # zero the parameter gradients                optimizer.zero_grad()                # track history if only in train                with torch.set_grad_enabled(phase == 'Train'):                    outputs = model(inputs)                    loss = criterion(outputs['out'], masks)                    y_pred = outputs['out'].data.cpu().numpy().ravel()                    y_true = masks.data.cpu().numpy().ravel()                    for name, metric in metrics.items():                        if name == 'f1_score':                            # Use a classification threshold of 0.1                            batchsummary[f'{phase}_{name}'].append(                                metric(y_true > 0, y_pred > 0.1))                        else:                            batchsummary[f'{phase}_{name}'].append(                                metric(y_true.astype('uint8'), y_pred))                    # backward + optimize only if in training phase                    if phase == 'Train':                        loss.backward()                        optimizer.step()            batchsummary['epoch'] = epoch            epoch_loss = loss            batchsummary[f'{phase}_loss'] = epoch_loss.item()            print('{} Loss: {:.4f}'.format(phase, loss))        for field in fieldnames[3:]:            batchsummary[field] = np.mean(batchsummary[field])        print(batchsummary)        with open(os.path.join(bpath, 'log.csv'), 'a', newline='') as csvfile:            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)            writer.writerow(batchsummary)            # deep copy the model            if phase == 'Test' and loss < best_loss:                best_loss = loss                best_model_wts = copy.deepcopy(model.state_dict())    time_elapsed = time.time() - since    print('Training complete in {:.0f}m {:.0f}s'.format(        time_elapsed // 60, time_elapsed % 60))    print('Lowest Loss: {:4f}'.format(best_loss))    # load best model weights    model.load_state_dict(best_model_wts)    return model

確保將模型以及輸入和標籤傳送到同一裝置(可以是cpu或cuda)。

在進行正向和反向傳播以及引數更新之前,請記住使用optimizer.zero_grad()清除梯度。

訓練時,使用mode.train()將模型設定為訓練模式

進行推斷時,請使用mode.eval()將模型設定為評估模式。 這一點非常重要,因為這可以確保調整網路引數,以解決影響網路權重的批處理規範,丟失等技術。

最佳模型取決於最低的損失值。 您也可以根據評估指標選擇最佳模型。 但是您必須稍微修改一下程式碼。

我已使用均方誤差(MSE)損失函式完成此任務。 我使用MSE的原因是它是一個簡單的函式,可以提供更好的結果,並且可以為計算梯度提供更好的表面。 在我們的案例中,損失是在畫素級別上計算的,定義如下:

為了評估模型的定量效能,選擇了兩個評估指標。 第一個指標是受試者工作特徵曲線(ROC)和曲線下面積(AUC)測量[8]。 AUC或ROC是任何二元分類器(在這種情況下為二元分割掩碼)的程度或可分離性的可靠度量。 它提供了所有可能的分類閾值下模型效能的彙總度量。 優秀的模型具有接近於AUROC的值,這意味著分類器實際上與特定閾值的選擇無關。 用於評估的第二個指標是F1分數。 它定義為精度(P)和召回率(R)的諧波平均值,由以下方程式給出。

F1分數在1時達到最高值,在0時達到最差值。 對於分類任務,這是一個可靠的選擇,因為它同時考慮了誤報。

結果

最佳模型的測試AUROC值為0.842。 這是一個很高的分數,也反映在閾值操作之後獲得的分段輸出中。

下圖顯示了訓練期間的損失和評估指標。

我們可以觀察到,在整個訓練過程中,損失值逐漸減小。AUROC和F1評分隨著訓練的進行而提高。然而,我們看到無論是訓練還是驗證,F1的得分值都始終較低。事實上,這些都是糟糕的表現。產生這樣結果的原因是我在計算這個度量時使用了0.1的閾值。這不是基於資料集選擇的。F1分數值可以根據閾值的不同而變化。然而,AUROC是一個考慮了所有可能的閾值的健壯度量。因此,當您有一個二元分類任務時,使用AUROC度量是明智的。儘管模型在資料集上表現良好,從分割輸出影象中可以看出,與地面真實值相比,掩模被過度放大了。也許因為模型比需要的更深,我們正在觀察這種行為。如果你對此現象有任何評論,請發表評論,我想知道你的想法。

總結

我們學習瞭如何使用PyTorch中的DeepLabv3對我們的自定義資料集進行語義分割任務的遷移學習。

首先,我們瞭解了影象分割和遷移學習。

接下來,我們瞭解瞭如何建立用於分割的資料集類來訓練模型。

接下來是如何根據我們的資料集改變DeepLabv3模型的分割頭的最重要的一步。

在CrackForest資料集上對該方法進行了道路裂縫檢測測試。在僅僅經歷了25個時代之後,它的AUROC評分就達到了0.842。

程式碼可以在https://github.com/msminhas93/DeepLabv3FineTuning上找到。

感謝你閱讀這篇文章。希望你能從這篇文章中學到一些新的東西。

引用

[1] Rethinking Atrous Convolution for Semantic Image Segmentation, arXiv:1706.05587, Available: https://arxiv.org/abs/1706.05587

[2] Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation, arXiv:1802.02611, Available: https://arxiv.org/abs/1802.02611

[3] https://scikit-image.org/docs/dev/userguide/tutorialsegmentation.html

[4] Anomaly Detection in Images, arXiv:1905.13147, Available: https://arxiv.org/abs/1905.13147

[5] Yong Shi, Limeng Cui, Zhiquan Qi, Fan Meng, and Zhensong Chen. Automatic road crack detection using randomstructured forests.IEEE Transactions on Intelligent Transportation Systems, 17(12):3434–3445, 2016.

[6] https://github.com/cuilimeng/CrackForest-dataset

[7] AnoNet: Weakly Supervised Anomaly Detection in Textured Surfaces, arXiv:1911.10608, Available: https://arxiv.org/abs/1911.10608

[8] Charles X. Ling, Jin Huang, and Harry Zhang. Auc: A statistically consistent and more discriminating measurethan accuracy. InProceedings of the 18th International Joint Conference on Artificial Intelligence, IJCAI'03,pages 519–524, San Francisco, CA, USA, 2003. Morgan Kaufmann Publishers Inc.

25
最新評論
  • BSA-TRITC(10mg/ml) TRITC-BSA 牛血清白蛋白改性標記羅丹明
  • 全鏈路智慧測試及錄製回放最佳實踐