當我在使用深度學習進行圖像語義分割並想使用PyTorch在DeepLabv3[1]上運行一些實驗時,我找不到任何在線教程。並且torchvision不僅沒有提供分割數據集,而且也沒有關於DeepLabv3類內部結構的詳細解釋。然而,我是通過自己的研究進行了現有模型的遷移學習,我想分享這個過程,這樣可能會對你們有幫助。
在本文中,我將介紹如何使用預先訓練的語義分割DeepLabv3模型,通過使用遷移學習在PyTorch中進行道路裂縫檢測。同樣的過程也可以應用於調整自定義數據集的網絡。
介紹
讓我們首先簡要介紹圖像分割。 分割任務的主要目標是輸出像素級輸出蒙版,其中將屬於某些類別的區域分配給相同的不同像素值。 如果通過為每個類別分配不同的顏色來對這些細分蒙版進行顏色編碼以使其可視化,那麼你就會得到一個類似於兒童塗色書中的圖像。下面顯示了一個示例:
分割在計算機視覺和圖像處理領域已經存在很長時間了。 其中一些技術是簡單的閾值化,基於聚類的方法,例如k均值聚類分割,區域增長方法等。[3]
隨著深度學習的最新進展以及卷積神經網絡在圖像相關任務中比傳統方法的成功,這些技術也已應用於圖像分割任務。
這些網絡架構之一是Google的DeepLabv3。 對模型的工作原理進行解釋超出了本文的範圍。 相反,我們將專注於如何對數據集使用經過預訓練的DeepLabv3網絡。 為此,我們將簡要討論轉移學習。
遷移學習
當有限的數據可用時,深度學習模型往往會遇到困難。 對於大多數實際應用,即使不是不可能,也很難訪問大量數據集。 標註既繁瑣又費時。 即使您打算將其外包,您仍然必須花錢。 已經做出努力以能夠從有限的數據訓練模型。 這些技術中的一種稱為轉移學習。
遷移學習涉及使用針對源域和任務進行預訓練的網絡(希望您可以在其中訪問大型數據集),並將其用於您的預期/目標域和任務(與原始任務和域類似))[4]。下圖可以從概念上表示它。
我們根據自己的要求更改目標細分子網絡,然後訓練部分網絡或整個網絡。選擇的學習率低於正常訓練的學習率。 這是因為網絡已經對源任務具有良好的權重。 我們不想太快地改變權重。有時也可以凍結初始層,因為有人認為這些層提取了一般特徵,可以潛在地使用而無需任何更改。
接下來,在繼續PyTorch相關部分之前,我將討論本文中使用的數據集。
使用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:此方法是數據集對象將初始化的位置。 通常,您需要構建圖像文件路徑和相應的標籤,它們是用於分割的遮罩文件路徑。 然後,在len和getitem方法中使用這些路徑。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.
作者:Manpreet Singh Minhas
deephub翻譯組