DeepLabv3模型
Torchvision有可用的预训练模型,我们将使用其中一种模型。我编写了以下函数,该函数为您提供了具有自定义数量的输出通道的模型。如果您有多个班级,则可以更改此值。
""" DeepLabv3 Model download and change the head for your prediction"""fromtorchvision.models.segmentation.deeplabv3importDeepLabHeadfromtorchvisionimportmodelsdefcreateDeepLabv3(outputchannels=1): """DeepLabv3 class with custom headArgs:outputchannels (int, optional): The number of output channelsin 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) #Setthemodelintrainingmodemodel.train() returnmodel
首先,我们使用models.segmentation.deeplabv3_resnet101方法获得预训练模型,该方法将预训练模型下载到我们的系统缓存中。注意resnet101是从此特定方法获得的deeplabv3模型的基础模型。这决定了传递到分类器的特征向量的长度。
第二步是修改分割头即分类器的主要步骤。该分类器是网络的一部分,负责创建最终的细分输出。通过用具有新数量的输出通道的新DeepLabHead替换模型的分类器模块来完成更改。resnet101主干的特征向量大小为2048。如果您决定使用另一个主干,请相应地更改此值。
最后,我们将模型设置为训练模式。此步骤是可选的,因为您也可以在训练逻辑中执行此操作。
下一步是训练模型。
我定义了以下训练模型的train_model函数。它将训练和验证损失以及指标(如果指定)值保存到CSV日志文件中,以便于访问。训练代码代码如下。后面有充分的文档来解释发生了什么。
deftrain_model(model, criterion, dataloaders, optimizer, metrics, bpath, num_epochs): since=time.time() best_model_wts=copy.deepcopy(model.state_dict()) best_loss=1e10#Usegpuifavailabledevice=torch.device("cuda:0"iftorch.cuda.is_available() else"cpu") model.to(device) #Initializethelogfilefortrainingandtestinglossandmetricsfieldnames= ['epoch', 'Train_loss', 'Test_loss'] +\ [f'Train_{m}'forminmetrics.keys()] +\ [f'Test_{m}'forminmetrics.keys()] withopen(os.path.join(bpath, 'log.csv'), 'w', newline='') ascsvfile: writer=csv.DictWriter(csvfile, fieldnames=fieldnames) writer.writeheader() forepochinrange(1, num_epochs+1): print('Epoch {}/{}'.format(epoch, num_epochs)) print('-'*10) #Eachepochhasatrainingandvalidationphase#Initializebatchsummarybatchsummary= {a: [0] forainfieldnames} forphasein ['Train', 'Test']: ifphase=='Train': model.train() #Setmodeltotrainingmodeelse: model.eval() #Setmodeltoevaluatemode#Iterateoverdata. forsampleintqdm(iter(dataloaders[phase])): inputs=sample['image'].to(device) masks=sample['mask'].to(device) #zerotheparametergradientsoptimizer.zero_grad() #trackhistoryifonlyintrainwithtorch.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() forname, metricinmetrics.items(): ifname=='f1_score': #Useaclassificationthresholdof0.1batchsummary[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+optimizeonlyifintrainingphaseifphase=='Train': loss.backward() optimizer.step() batchsummary['epoch'] =epochepoch_loss=lossbatchsummary[f'{phase}_loss'] =epoch_loss.item() print('{} Loss: {:.4f}'.format(phase, loss)) forfieldinfieldnames[3:]: batchsummary[field] =np.mean(batchsummary[field]) print(batchsummary) withopen(os.path.join(bpath, 'log.csv'), 'a', newline='') ascsvfile: writer=csv.DictWriter(csvfile, fieldnames=fieldnames) writer.writerow(batchsummary) #deepcopythemodelifphase=='Test'andloss<best_loss: best_loss=lossbest_model_wts=copy.deepcopy(model.state_dict()) time_elapsed=time.time() -sinceprint('Training complete in {:.0f}m {:.0f}s'.format( time_elapsed// 60, time_elapsed % 60))print('Lowest Loss: {:4f}'.format(best_loss)) #loadbestmodelweightsmodel.load_state_dict(best_model_wts) returnmodel
确保将模型以及输入和标签发送到同一设备(可以是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/user_guide/tutorial_segmentation.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.



