torchvision.models.segmentation.deeplabv3_resnet50实战

最近在进行医学图像处理时，使用到了预训练好的deeplabv3-resnet50网络，在此记录一下应用（调参）的过程。

首先定义模型类，由于直接print网络结构可知，torchvision导入的预训练的模型最后一层是一个conv层，所以需要我们增添一个激活层。【注：导入模型的输出为字典格式，存在‘out’中】

由于最终输出结果为灰度图，是一种二分类逐像素预测，所以此处num_classes设置为2.

class seg_(torch.nn.Module):
    def __init__(self):
        super(seg_,self).__init__()
        self.model1 = torchvision.models.segmentation.deeplabv3_resnet50(pretrained = False,num_classes=2,pretrained_backbone=False)
        self.act = torch.nn.Sigmoid()
    def forward(self,input):
        x = self.model1(input)['out']
        x = self.act(x)
        return x

trainset,testset = make_datasets('mix_datasets/','mix_mask/')

train_loader = torch.utils.data.DataLoader(dataset = trainset,
                                          batch_size = batch_size,
                                          shuffle = True,
                                          num_workers = 1,
                                          )
val_loader = torch.utils.data.DataLoader(dataset = testset,
                                        batch_size = batch_size,
                                        shuffle = False,
                                        num_workers = 1)

使用自己的make_datasets函数获取训练集以及测试集

 实例化模型，定义优化器和损失函数，此处应用的是交叉熵损失函数，最好是可以增添一个DICE_LOSS损失函数。

import torch.optim as optim
seg_model = seg_().cuda()
optimizer = optim.SGD(seg_model.parameters(),lr=0.001,momentum=0.9,weight_decay=1e-4,nesterov=True)
loss_fn = torch.nn.CrossEntropyLoss()

 接下来定义训练函数，label传入为（batch_size,1,h,w），需要将其squeeze（）掉channel通道，才能进行交叉熵损失函数计算。

验证集使用的是交并比（MIOU）来进行模型准确率的衡量。

最终将训练过程以图像的形式保存在文件夹中。

def train(num_epochs = num_epoches, model = seg_model, optimizer = optimizer, 
          loss_fn = loss_fn, train_loader = train_loader,
         test_loader = val_loader):
    epoches_list = [i for i in range(num_epochs)]
    epoch_losses_train = []
    epoch_losses_val = []
    acc_list = []
    for epoch in range(num_epochs):
        print('epoch: %d/%d' % (epoch+1,num_epochs))
        model.train()
        batch_losses = []
        corr = 0
        total = 0
        for step,(imgs,label) in enumerate(tqdm(train_loader)):
            imgs = imgs.cuda() #(batch_size,3,224,224)
            label = label.type(torch.LongTensor).cuda()#(batch_size,1,224,224)
            label = label.squeeze()#(batch_size,224,224)
            outputs = model(imgs)#(batch_size,2,224,224)

            loss = loss_fn(outputs,label) #交叉熵损失函数，最好可以加上Dice_Loss()
            loss_value = loss.data.cpu().numpy() #在转为numpy之前先从gpu转为cpu

            batch_losses.append(loss_value)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            t.set_description('epoch: %i'%epoch)

        epoch_loss = np.mean(batch_losses)
        epoch_losses_train.append(epoch_loss)
        print ("train loss: %g" % epoch_loss)
        
        print('########start val#########')
        model.eval()
        batch_losses = []
        for step, (imgs, label_imgs) in enumerate(tqdm(test_loader)):
            with torch.no_grad(): 
                imgs = imgs.cuda() 
                label_imgs = label_imgs.type(torch.LongTensor).cuda()
                label_imgs = label_imgs.squeeze()
            
                outputs = model(imgs)
                _,pred = outputs.max(dim=1)
                corr += compute_iou(pred.cpu().numpy(),label_imgs.cpu().numpy(),0.5)
                total += label_imgs.size(0)
                # compute the loss:
                outputs=outputs[:label_imgs.size(0)]
                loss = loss_fn(outputs, label_imgs)
                
                loss_value = loss.data.cpu().numpy()
                batch_losses.append(loss_value)
                bar.set_description("val loss: {:3f} val acc{:3f}%".format(epoch_loss,corr/total*100))
        epoch_loss = np.mean(batch_losses)
        epoch_losses_val.append(epoch_loss)
        acc_list.append(corr/total*100)
        
         # save the model weights to disk:
        checkpoint_path = "./seg_model_deeplab50/deeplab50_segmentation_model_false"+"_epoch_" + str(epoch+1) + ".pth"
        torch.save(model.state_dict(), checkpoint_path)
    
    #生成图片
    # 设置图片大小和分辨率
    fig=plt.figure(figsize=(8, 6), dpi=160)
    # 绘制折线图
    plt.plot(epoches_list, epoch_losses_train, label='Loss',color='r')
    # 描述内容信息
    plt.xlabel("epochs")  # x轴
    plt.ylabel("Train Loss")  # y轴
    plt.title("Loss")
    plt.savefig("./seg_model_deeplab50_false_Loss.png")  # 保存在当前目录中

    # 设置图片大小和分辨率
    fig=plt.figure(figsize=(8, 6), dpi=160)
    #设置横坐标范围，这里要与上面a1中的数据个数相对应
    #plt.ylim([0.2, 7])#设置y轴显示范围从0.8到1
    # 绘制折线图
    plt.plot(epoches_list, acc_list, label='Accuracy',color='b')
    # 描述内容信息
    plt.xlabel("epochs")  # x轴
    plt.ylabel("Accuracy")  # y轴
    plt.title("Accuracy")
    plt.savefig("./seg_model_deeplab50_false_Acc.png")  # 保存在当前目录中 

 交并比函数如下。

def _fast_hist(row_label, row_image, n_class):
    mask = (row_label>= 0) & (row_label < n_class)
    hist = np.bincount(
        n_class * row_label[mask] +
        row_image[mask], minlength=n_class ** 2).reshape(n_class, n_class)
    return hist

def compute_iou(result, target,threshold=0.5):
    #iou阈值默认为50%
    correct = 0
    for single_image,single_label in zip(result,target):
        hist = 0
        for row_image,row_label in zip(single_image,single_label):
            hist += _fast_hist(row_label.flatten(), row_image.flatten(), 2)
        iu = np.diag(hist) / (0.000001+hist.sum(axis=1) + hist.sum(axis=0) - np.diag(hist))
        if np.mean(iu)>=threshold:
            correct += 1
    return correct

欢迎大家交流。