语义分割是计算机视觉的一项重要任务，本教程使用Jittor框架实现了DeepLabV3+语义分割模型。

DeepLabV3+论文：https://arxiv.org/pdf/1802.02611.pdf

完整代码：https://github.com/Jittor/deeplab-jittor

1. 数据集

1.1 数据准备

VOC2012数据集是目标检测、语义分割等任务常用的数据集之一，本教程使用VOC数据集的2012 trainaug (train + sbd set)作为训练集，2012 val set作为测试集。

VOC数据集中的物体共包括20个前景类别：'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor' 和背景类别

示例4：语义分割之 DeepLabV3+ - 图1

最终数据集的文件组织如下。

  1. # 文件组织
  2. 根目录
  3. |----voc_aug
  4. |    |----datalist
  5. |    |    |----train.txt
  6. |    |    |----val.txt
  7. |    |----images
  8. |    |----annotations

1.2 数据加载

使用jittor.dataset.dataset的基类Dataset可以构造自己的数据集，需要实现init、getitem、函数。

  • init: 定义数据路径，这里的data_root需设置为之前您设定的 voc_aug, split 为 train val test 之一，表示选择训练集、验证集还是测试集。同时需要调用self.set_attr来指定数据集加载所需的参数batch_size，total_len、shuffle。
  • getitem: 返回单个item的数据。
  1. import numpy as np
  2. import os
  3. from PIL import Image
  4. import matplotlib.pyplot as plt
  5. from jittor.dataset.dataset import Dataset, dataset_root
  6. import jittor as jt
  7. import os
  8. import os.path as osp
  9. from PIL import Image, ImageOps, ImageFilter
  10. import numpy as np
  11. import scipy.io as sio
  12. import random
  14. def fetch(image_path, label_path):
  15.     with open(image_path, 'rb') as fp:
  16.         image = Image.open(fp).convert('RGB')
  18.     with open(label_path, 'rb') as fp:
  19.         label = Image.open(fp).convert('P')
  21.     return image, label
  24. def scale(image, label):
  25.     SCALES = (0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0)
  26.     ratio = np.random.choice(SCALES)
  27.     w,h = image.size
  28.     nw = (int)(w*ratio)
  29.     nh = (int)(h*ratio)
  31.     image = image.resize((nw, nh), Image.BILINEAR)
  32.     label = label.resize((nw, nh), Image.NEAREST)
  34.     return image, label
  37. def pad(image, label):
  38.     w,h = image.size
  39.     crop_size = 513
  40.     pad_h = max(crop_size - h, 0)
  41.     pad_w = max(crop_size - w, 0)
  42.     image = ImageOps.expand(image, border=(0, 0, pad_w, pad_h), fill=0)
  43.     label = ImageOps.expand(label, border=(0, 0, pad_w, pad_h), fill=255)
  45.     return image, label
  48. def crop(image, label):
  49.     w, h = image.size
  50.     crop_size = 513
  51.     x1 = random.randint(0, w - crop_size)
  52.     y1 = random.randint(0, h - crop_size)
  53.     image = image.crop((x1, y1, x1 + crop_size, y1 + crop_size))
  54.     label = label.crop((x1, y1, x1 + crop_size, y1 + crop_size))
  57.     return image, label
  60. def normalize(image, label):
  61.     mean = (0.485, 0.456, 0.40)
  62.     std = (0.229, 0.224, 0.225)
  63.     image = np.array(image).astype(np.float32)
  64.     label = np.array(label).astype(np.float32)
  66.     image /= 255.0
  67.     image -= mean
  68.     image /= std
  69.     return image, label
  72. def flip(image, label):
  73.     if random.random() < 0.5:
  74.         image = image.transpose(Image.FLIP_LEFT_RIGHT)
  75.         label = label.transpose(Image.FLIP_LEFT_RIGHT)
  76.     return image, label
  79. class BaseDataset(Dataset):
  80.     def __init__(self,  data_root='/voc/', split='train', batch_size=1, shuffle=False):
  81.         super().__init__()
  82.         ''' total_len , batch_size, shuffle must be set '''
  83.         self.data_root = data_root
  84.         self.split = split
  85.         self.batch_size = batch_size
  86.         self.shuffle = shuffle
  88.         self.image_root = os.path.join(data_root, 'images')
  89.         self.label_root = os.path.join(data_root, 'annotations')
  91.         self.data_list_path = os.path.join(self.data_root,'/datalist/' + self.split + '.txt')
  92.         self.image_path = []
  93.         self.label_path = []
  95.         with open(self.data_list_path, "r") as f:
  96.             lines = f.read().splitlines()
  98.         for idx, line in enumerate(lines):
  99.             _img_path = os.path.join(self.image_root, line + '.jpg')
  100.             _label_path = os.path.join(self.label_root, line + '.png')
  102.             assert os.path.isfile(_img_path)
  103.             assert os.path.isfile(_label_path)
  104.             self.image_path.append(_img_path)
  105.             self.label_path.append(_label_path)
  106.         self.total_len = len(self.image_path)
  108.         # set_attr must be called to set batch size total len and shuffle like __len__ function in pytorch
  109.         self.set_attr(batch_size = self.batch_size, total_len = self.total_len, shuffle = self.shuffle) # bs , total_len, shuffle
  112.     def __getitem__(self, image_id):
  113.         return NotImplementedError
  116. class TrainDataset(BaseDataset):
  117.     def __init__(self,  data_root='/voc/', split='train', batch_size=1, shuffle=False):
  118.         super(TrainDataset, self).__init__(data_root, split, batch_size, shuffle)
  120.     def __getitem__(self, image_id):
  121.         image_path = self.image_path[image_id]
  122.         label_path = self.label_path[image_id]
  123.         image, label = fetch(image_path, label_path)
  124.         image, label = scale(image, label)
  125.         image, label = pad(image, label)
  126.         image, label = crop(image, label)
  127.         image, label = flip(image, label)
  128.         image, label = normalize(image, label)
  129.         image = np.array(image).astype(np.float).transpose(2, 0, 1)
  130.         image = jt.array(image)
  131.         label = jt.array(np.array(label).astype(np.int))
  132.         return image, label
  135. class ValDataset(BaseDataset):
  136.     def __init__(self,  data_root='/voc/', split='train', batch_size=1, shuffle=False):
  137.         super(ValDataset, self).__init__(data_root, split, batch_size, shuffle)
  139.     def __getitem__(self, image_id):
  140.         image_path = self.image_path[image_id]
  141.         label_path = self.label_path[image_id]
  143.         image, label = fetch(image_path, label_path)
  144.         image, label = normalize(image, label)
  146.         image = np.array(image).astype(np.float).transpose(2, 0, 1)
  147.         image = jt.array(image)
  148.         label = jt.array(np.array(label).astype(np.int))
  150.         return image, label

2. 模型定义

示例4：语义分割之 DeepLabV3+ - 图2

上图为DeepLabV3+论文给出的网络架构图。本教程采用ResNe为backbone。输入图像尺寸为513*513。

整个网络可以分成 backbone aspp decoder 三个部分。


2.1 backbonb 这里使用最常见的 ResNet 作为 backbone 并且在ResNet的最后两次使用空洞卷积来扩大感受野，其完整定义如下：


  1. import jittor as jt
  2. from jittor import nn
  3. from jittor import Module
  4. from jittor import init
  5. from jittor.contrib import concat, argmax_pool
  6. import time
  9. class Bottleneck(Module):
  10.     expansion = 4
  11.     def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None):
  12.         super(Bottleneck, self).__init__()
  13.         self.conv1 = nn.Conv(inplanes, planes, kernel_size=1, bias=False)
  14.         self.bn1 = nn.BatchNorm(planes)
  15.         self.conv2 = nn.Conv(planes, planes, kernel_size=3, stride=stride,
  16.                                dilation=dilation, padding=dilation, bias=False)
  17.         self.bn2 = nn.BatchNorm(planes)
  18.         self.conv3 = nn.Conv(planes, planes * 4, kernel_size=1, bias=False)
  19.         self.bn3 = nn.BatchNorm(planes * 4)
  20.         self.relu = nn.ReLU()
  21.         self.downsample = downsample
  22.         self.stride = stride
  23.         self.dilation = dilation
  25.     def execute(self, x):
  26.         residual = x
  28.         out = self.conv1(x)
  29.         out = self.bn1(out)
  30.         out = self.relu(out)
  32.         out = self.conv2(out)
  33.         out = self.bn2(out)
  34.         out = self.relu(out)
  36.         out = self.conv3(out)
  37.         out = self.bn3(out)
  39.         if self.downsample is not None:
  40.             residual = self.downsample(x)
  42.         out += residual
  43.         out = self.relu(out)
  45.         return out
  48. class ResNet(Module):
  49.     def __init__(self, block, layers, output_stride):
  50.         super(ResNet, self).__init__()
  51.         self.inplanes = 64
  52.         blocks = [1, 2, 4]
  53.         if output_stride == 16:
  54.             strides = [1, 2, 2, 1]
  55.             dilations = [1, 1, 1, 2]
  56.         elif output_stride == 8:
  57.             strides = [1, 2, 1, 1]
  58.             dilations = [1, 1, 2, 4]
  59.         else:
  60.             raise NotImplementedError
  62.         # Modules
  63.         self.conv1 = nn.Conv(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
  64.         self.bn1 = nn.BatchNorm(64)
  65.         self.relu = nn.ReLU()
  66.         # self.maxpool = nn.Pool(kernel_size=3, stride=2, padding=1)
  68.         self.layer1 = self._make_layer(block, 64, layers[0], stride=strides[0], dilation=dilations[0])
  69.         self.layer2 = self._make_layer(block, 128, layers[1], stride=strides[1], dilation=dilations[1])
  70.         self.layer3 = self._make_layer(block, 256, layers[2], stride=strides[2], dilation=dilations[2])
  71.         self.layer4 = self._make_MG_unit(block, 512, blocks=blocks, stride=strides[3], dilation=dilations[3])
  74.     def _make_layer(self, block, planes, blocks, stride=1, dilation=1):
  75.         downsample = None
  76.         if stride != 1 or self.inplanes != planes * block.expansion:
  77.             downsample = nn.Sequential(
  78.                 nn.Conv(self.inplanes, planes * block.expansion,
  79.                           kernel_size=1, stride=stride, bias=False),
  80.                 nn.BatchNorm(planes * block.expansion),
  81.             )
  83.         layers = []
  84.         layers.append(block(self.inplanes, planes, stride, dilation, downsample))
  85.         self.inplanes = planes * block.expansion
  86.         for i in range(1, blocks):
  87.             layers.append(block(self.inplanes, planes, dilation=dilation))
  89.         return nn.Sequential(*layers)
  91.     def _make_MG_unit(self, block, planes, blocks, stride=1, dilation=1):
  92.         downsample = None
  93.         if stride != 1 or self.inplanes != planes * block.expansion:
  94.             downsample = nn.Sequential(
  95.                 nn.Conv(self.inplanes, planes * block.expansion,
  96.                           kernel_size=1, stride=stride, bias=False),
  97.                 nn.BatchNorm(planes * block.expansion),
  98.             )
  100.         layers = []
  101.         layers.append(block(self.inplanes, planes, stride, dilation=blocks[0]*dilation,
  102.                             downsample=downsample))
  103.         self.inplanes = planes * block.expansion
  104.         for i in range(1, len(blocks)):
  105.             layers.append(block(self.inplanes, planes, stride=1,
  106.                                 dilation=blocks[i]*dilation))
  108.         return nn.Sequential(*layers)
  110.     def execute(self, input):
  112.         x = self.conv1(input)
  113.         x = self.bn1(x)
  114.         x = self.relu(x)
  115.         x = argmax_pool(x, 2, 2)
  116.         x = self.layer1(x)
  118.         low_level_feat = x
  119.         x = self.layer2(x)
  120.         x = self.layer3(x)
  122.         x = self.layer4(x)
  123.         return x, low_level_feat
  125. def resnet50(output_stride):
  126.     model = ResNet(Bottleneck, [3,4,6,3], output_stride)
  127.     return model
  129. def resnet101(output_stride):
  130.     model = ResNet(Bottleneck, [3,4,23,3], output_stride)
  131.     return model

2.2 ASPP 即使用不同尺寸的 dilation conv 对 backbone 得到的 feature map 进行卷积，最后 concat 并整合得到新的特征。


  1. import jittor as jt
  2. from jittor import nn
  3. from jittor import Module
  4. from jittor import init
  5. from jittor.contrib import concat
  8. class Single_ASPPModule(Module):
  9.     def __init__(self, inplanes, planes, kernel_size, padding, dilation):
  10.         super(Single_ASPPModule, self).__init__()
  11.         self.atrous_conv = nn.Conv(inplanes, planes, kernel_size=kernel_size,
  12.                                             stride=1, padding=padding, dilation=dilation, bias=False)
  13.         self.bn = nn.BatchNorm(planes)
  14.         self.relu = nn.ReLU()
  16.     def execute(self, x):
  17.         x = self.atrous_conv(x)
  18.         x = self.bn(x)
  19.         x = self.relu(x)
  20.         return x
  22. class ASPP(Module):
  23.     def __init__(self, output_stride):
  24.         super(ASPP, self).__init__()
  25.         inplanes = 2048
  26.         if output_stride == 16:
  27.             dilations = [1, 6, 12, 18]
  28.         elif output_stride == 8:
  29.             dilations = [1, 12, 24, 36]
  30.         else:
  31.             raise NotImplementedError
  33.         self.aspp1 = Single_ASPPModule(inplanes, 256, 1, padding=0, dilation=dilations[0])
  34.         self.aspp2 = Single_ASPPModule(inplanes, 256, 3, padding=dilations[1], dilation=dilations[1])
  35.         self.aspp3 = Single_ASPPModule(inplanes, 256, 3, padding=dilations[2], dilation=dilations[2])
  36.         self.aspp4 = Single_ASPPModule(inplanes, 256, 3, padding=dilations[3], dilation=dilations[3])
  37.         self.global_avg_pool = nn.Sequential(GlobalPooling(),
  38.                                              nn.Conv(inplanes, 256, 1, stride=1, bias=False),
  39.                                              nn.BatchNorm(256),
  40.                                              nn.ReLU())
  41.         self.conv1 = nn.Conv(1280, 256, 1, bias=False)
  43.         self.bn1 = nn.BatchNorm(256)
  44.         self.relu = nn.ReLU()
  45.         self.dropout = nn.Dropout(0.5)
  47.     def execute(self, x):
  48.         x1 = self.aspp1(x)
  49.         x2 = self.aspp2(x)
  50.         x3 = self.aspp3(x)
  51.         x4 = self.aspp4(x)
  52.         x5 = self.global_avg_pool(x)
  53.         x5 = x5.broadcast((1,1,x4.shape[2],x4.shape[3]))
  54.         x = concat((x1, x2, x3, x4, x5), dim=1)
  55.         x = self.conv1(x)
  56.         x = self.bn1(x)
  57.         x = self.relu(x)
  58.         x = self.dropout(x)
  59.         return x
  61. class GlobalPooling (Module):
  62.     def __init__(self):
  63.         super(GlobalPooling, self).__init__()
  64.     def execute (self, x):
  65.         return jt.mean(x, dims=[2,3], keepdims=1)

2.3 Decoder: Decoder 将 ASPP 的特征放大后与 ResNet 的中间特征一起 concat 得到最后分割所用的特征。


  1. import jittor as jt
  2. from jittor import nn
  3. from jittor import Module
  4. from jittor import init
  5. from jittor.contrib import concat
  6. import time
  8. class Decoder(nn.Module):
  9.     def __init__(self, num_classes):
  10.         super(Decoder, self).__init__()
  11.         low_level_inplanes = 256
  13.         self.conv1 = nn.Conv(low_level_inplanes, 48, 1, bias=False)
  14.         self.bn1 = nn.BatchNorm(48)
  15.         self.relu = nn.ReLU()
  16.         self.last_conv = nn.Sequential(nn.Conv(304, 256, kernel_size=3, stride=1, padding=1, bias=False),
  17.                                        nn.BatchNorm(256),
  18.                                        nn.ReLU(),
  19.                                        nn.Dropout(0.5),
  20.                                        nn.Conv(256, 256, kernel_size=3, stride=1, padding=1, bias=False),
  21.                                        nn.BatchNorm(256),
  22.                                        nn.ReLU(),
  23.                                        nn.Dropout(0.1),
  24.                                        nn.Conv(256, num_classes, kernel_size=1, stride=1, bias=True))
  26.     def execute(self, x, low_level_feat):
  27.         low_level_feat = self.conv1(low_level_feat)
  28.         low_level_feat = self.bn1(low_level_feat)
  29.         low_level_feat = self.relu(low_level_feat)
  31.         x_inter = nn.resize(x, size=(low_level_feat.shape[2], low_level_feat.shape[3]) , mode='bilinear')
  32.         x_concat = concat((x_inter, low_level_feat), dim=1)
  33.         x = self.last_conv(x_concat)
  34.         return x

2.4 完整的模型整合如下： 即将以上部分通过一个类连接起来。


  1. import jittor as jt
  2. from jittor import nn
  3. from jittor import Module
  4. from jittor import init
  5. from jittor.contrib import concat
  6. from decoder import Decoder
  7. from aspp import ASPP
  8. from backbone import resnet50, resnet101
  10. class DeepLab(Module):
  11.     def __init__(self, output_stride=16, num_classes=21):
  12.         super(DeepLab, self).__init__()
  13.         self.backbone = resnet101(output_stride=output_stride)
  14.         self.aspp = ASPP(output_stride)
  15.         self.decoder = Decoder(num_classes)
  17.     def execute(self, input):
  18.         x, low_level_feat = self.backbone(input)
  19.         x = self.aspp(x)
  20.         x = self.decoder(x, low_level_feat)
  21.         x = nn.resize(x, size=(input.shape[2], input.shape[3]), mode='bilinear')
  22.         return x

3. 模型训练


3.1 模型训练参数设定如下：


  1. # Learning parameters
  2. batch_size = 8
  3. learning_rate = 0.005
  4. momentum = 0.9
  5. weight_decay = 1e-4
  6. epochs = 50

3.2 定义模型、优化器、数据加载器。


  1. model = DeepLab(output_stride=16, num_classes=21)
  2. optimizer = nn.SGD(model.parameters(), 
  3.                    lr,
  4.                    momentum=momentum, 
  5.                    weight_decay=weight_decay)
  6. train_loader = TrainDataset(data_root='/vocdata/',
  7.                             split='train',
  8.                             batch_size=batch_size,
  9.                             shuffle=True)
  10. val_loader = ValDataset(data_root='/vocdata/',
  11.                         split='val',
  12.                         batch_size=1,
  13.                         shuffle=False)

3.3 模型训练与验证


  1. # lr scheduler
  2. def poly_lr_scheduler(opt, init_lr, iter, epoch, max_iter, max_epoch):
  3.     new_lr = init_lr * (1 - float(epoch * max_iter + iter) / (max_epoch * max_iter)) ** 0.9
  4.     opt.lr = new_lr
  6. # train function
  7. def train(model, train_loader, optimizer, epoch, init_lr):
  8.     model.train()
  9.     max_iter = len(train_loader)
  11.     for idx, (image, target) in enumerate(train_loader):
  12.         poly_lr_scheduler(optimizer, init_lr, idx, epoch, max_iter, 50) # using poly_lr_scheduler 
  13.         image = image.float32()
  14.         pred = model(image)
  15.         loss = nn.cross_entropy_loss(pred, target, ignore_index=255)
  16.         optimizer.step (loss)
  17.         print ('Training in epoch {} iteration {} loss = {}'.format(epoch, idx, loss.data[0]))
  19. # val function
  20. # we omit evaluator code and you can 
  21. def val (model, val_loader, epoch, evaluator):
  22.     model.eval()
  23.     evaluator.reset()
  24.     for idx, (image, target) in enumerate(val_loader):
  25.         image = image.float32()
  26.         output = model(image)
  27.         pred = output.data
  28.         target = target.data
  29.         pred = np.argmax(pred, axis=1)
  30.         evaluator.add_batch(target, pred)
  31.         print ('Test in epoch {} iteration {}'.format(epoch, idx))
  32.     Acc = evaluator.Pixel_Accuracy()
  33.     Acc_class = evaluator.Pixel_Accuracy_Class()
  34.     mIoU = evaluator.Mean_Intersection_over_Union()
  35.     FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
  36.     best_miou = 0.0
  38.     if (mIoU > best_miou):
  39.         best_miou = mIoU
  40.     print ('Testing result of epoch {} miou = {} Acc = {} Acc_class = {} \
  41.                 FWIoU = {} Best Miou = {}'.format(epoch, mIoU, Acc, Acc_class, FWIoU, best_miou))

3.4 evaluator 写法：使用混淆矩阵计算 Pixel accuracy 和 mIoU。


  1. class Evaluator(object):
  2.     def __init__(self, num_class):
  3.         self.num_class = num_class
  4.         self.confusion_matrix = np.zeros((self.num_class,)*2)
  6.     def Pixel_Accuracy(self):
  7.         Acc = np.diag(self.confusion_matrix).sum() / self.confusion_matrix.sum()
  8.         return Acc
  10.     def Pixel_Accuracy_Class(self):
  11.         Acc = np.diag(self.confusion_matrix) / self.confusion_matrix.sum(axis=1)
  12.         Acc = np.nanmean(Acc)
  13.         return Acc
  15.     def Mean_Intersection_over_Union(self):
  16.         MIoU = np.diag(self.confusion_matrix) / (
  17.                   np.sum(self.confusion_matrix, axis=1) + 
  18.                 np.sum(self.confusion_matrix, axis=0)-
  19.                 np.diag(self.confusion_matrix))
  20.         MIoU = np.nanmean(MIoU)
  21.         return MIoU
  23.     def Frequency_Weighted_Intersection_over_Union(self):
  24.         freq = np.sum(self.confusion_matrix, axis=1) / np.sum(self.confusion_matrix)
  25.         iu = np.diag(self.confusion_matrix) / (
  26.                     np.sum(self.confusion_matrix, axis=1) + np.sum(self.confusion_matrix, axis=0) -
  27.                     np.diag(self.confusion_matrix))
  29.         FWIoU = (freq[freq > 0] * iu[freq > 0]).sum()
  30.         return FWIoU
  32.     def _generate_matrix(self, gt_image, pre_image):
  33.         mask = (gt_image >= 0) & (gt_image < self.num_class)
  34.         label = self.num_class * gt_image[mask].astype('int') + pre_image[mask]
  35.         count = np.bincount(label, minlength=self.num_class**2)
  36.         confusion_matrix = count.reshape(self.num_class, self.num_class)
  37.         return confusion_matrix
  39.     def add_batch(self, gt_image, pre_image):
  40.         assert gt_image.shape == pre_image.shape
  41.         self.confusion_matrix += self._generate_matrix(gt_image, pre_image)
  43.     def reset(self):
  44.         self.confusion_matrix = np.zeros((self.num_class,) * 2)

3.5 训练入口函数


  1. epochs = 50
  2. evaluator = Evaluator(21)
  3. train_loader = TrainDataset(data_root='/voc/data/path/', split='train', batch_size=8, shuffle=True)
  4. val_loader = ValDataset(data_root='/voc/data/path/', split='val', batch_size=1, shuffle=False)
  5. learning_rate = 0.005
  6. momentum = 0.9
  7. weight_decay = 1e-4
  8. optimizer = nn.SGD(model.parameters(), learning_rate, momentum, weight_decay)
  10. for epoch in range (epochs):
  11.     train(model, train_loader, optimizer, epoch, learning_rate)
  12.     val(model, val_loader, epoch, evaluator)

4. 参考