DeepLab 是google 推出的用于语义分割的一系列论文，在语义分割领域取得 state of the art 表现。[1] [2] 语义分割排行榜 。
我的项目需要把人脸抠图出来从而去除背景的干扰，适合使用语义分割的方法，由于用自己写的 unet 训练后效果很差，又自我感觉数据集没什么问题，因而决定用 DeepLab 再试试。
Deeplab 架构及要点 这里以 DeepLabV3+[3]
要点可以总结为三部分：
左上 DCNN 即 feature extractor backbone，可选用 mobilenet_v2、resnet、xception等， 论文在这里贡献主要在于针对 xception 的改进，使其结构灵活化，尤其是通过 output_stride 设置，直接改变 encoder 向下潜的层数，方便使用者权衡 GPU 和网络效能。同时大量把空洞卷积应用在 seperable convolution 的 depthwize_conv 上，提高了效能并增大感受野，如下图：
Encoder 后半部分 ASPP(Atrous Spatial Pyramid Pooling)，结合了 SPP 的多尺度特性和空洞卷积的低参特性：[4]
Decoder 不是多层上采样，而是很简单暴力的只有两层快速上采样，这点让我比较意外。backbone 输出的 low-level feature 与 ASPP 的结果 concat 一起，而在 concat 前，使用 point_conv 来调整两者 channel 及比例，这个是比较影响结果的 trick，论文中有更多详细例证。
配置环境，跑通 deeplab 库的样例数据和模型。 Deeplab 由 Tensorflow 开源，属于 models/research/ 下的一个子项目，跟 object detection api 等项目一样依赖 slim， 因此需要 git clone 整个 models 项目，再根据教程 进行配置，以及测试基本数据集是否能正常训练。
按照这里的文档 尝试 PASCAL VOC 2012 数据集的训练，大概会下载2G数据，并随后转为 tfrecords。使用默认命令，尝试训练。model zoo 下载预训练模型，继续训练，跑一下，total loss 大概 0.2。这样我们对于训练方案和目标 loss 心里就大概有了底。
DeepLab 训练参数设置 DeepLab 训练对显存要求还是很高的，这里总结下应对显存少的办法：
按文档所说，从预训练模型开始训练，并设置 fine_tune_batch_norm = False，从而减少 bn 参数对 GPU 的占用。（当然，训练自己的数据集的话，这个方案可能不太可靠）
使用相对小一点的模型架构，默认 model_variant 参数为 xception_65，使用 mobilenet_v2 会小很多，可供选择的模型参考 model zoo 列表或者项目文件 core.feature_extractor.py 里声明的模型。
减小 crop_size 大小，原因待看论文补充，值不得小于 [321, 321] [5]
减少 batch_size， 但文档建议 batch_size 要足够大（>12）才能有较好效果。
output_stride 用 16。output_stride 是 encoder 输入到输出的分辨率比例[6]
我这里实测，在自己的单张 rtx2080 显卡上（win10），使用 mobilenet_v2， crop_size 513， batch_size 可以到 8，再大似乎就不行了。[7]
在单张 rtx2080ti 显卡上，使用 xception_65，不断压榨参数，最终能够在 449,449 crop_size 时用 batch_size=6 时跑起来，但是loss下降很慢，感觉就是因为 bn 没法充分训练，训练一晚，miou 还不到 0.7。 使用 mobilenet_v2，能够 batch_size=16 跑起来。
使用 mobilenet_v2 架构训练时，官方备注要把 atrous_stride 设为 None，但其实不必的，mobilenet_v2 也只是作为 backbone，设置了 ASPP 可以正常生效，官方是为了速度才不建议用[8]
准备数据 使用 CelebAMaskHQ 数据集 CelebAMaskHQ 提供了一个非常高质量的人脸图像数据库，由 CelebA 筛选而来，共 30000 个样本，图像分辨率为 1024x1024，语义分割标注分辨率 512x512，标注19个类别如眼睛、眼镜、鼻子等等。
由于我们的目的是去掉人像的背景，只要把整个人像标注为一个类而其他标注为背景就可以了，所以这些类别我们在实际使用时都会被看作一个类。原数据集标注数据是按照每个类别分别生成掩码的，我们需要预处理把他们汇总在一起，做个mask叠加的动作。预处理代码 ，或者参考下面我的脚本。（以下代码生成汇总mask(19分类)、人像mask(2分类)以及去掉背景的人像图）
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 import osimport sysimport cv2import globimport numpy as npimport threadingimport shutilfrom tqdm import tqdm_notebook,tnrangedef touch_dir (dir) : if not os.path.exists(dir): print('making dir : ' , dir) os.makedirs(dir) def clean_and_touch_dir (dir) : if os.path.exists(dir): shutil.rmtree(dir) os.makedirs(dir) def work_on_imgs (num_from, num_to) : for k in tnrange(num_from, num_to, desc='img {:d} to {:d}' .format(num_from, num_to)): folder_num = int(k / 2000 ) num_str = '{:05d}' .format(k) img_path = os.path.join(img_dir, str(k) + '.jpg' ) image = cv2.imread(img_path) image = cv2.resize(image, (512 , 512 ), interpolation=cv2.INTER_CUBIC) mask = np.zeros((512 , 512 ), dtype='uint8' ) for idx, label in enumerate(label_list): filename = os.path.join(mask_dir, str(folder_num), num_str + '_' + label + '.png' ) if (os.path.exists(filename)): im=cv2.imread(filename) im = im[:, :, 0 ] for i in range(512 ): for j in range(512 ): if im[i][j] != 0 : mask[i][j] = (idx + 1 ) total_mask_path = os.path.join(total_mask_dir, str(k) + '.png' ) cv2.imwrite(total_mask_path, mask) one_mask = np.uint8(mask>0 ) inverse_mask = 1 -one_mask bg = np.ones((512 , 512 , 3 ), dtype='uint8' ) * 255 masked = cv2.bitwise_and(image, image, mask=one_mask) bg_masked = cv2.bitwise_and(bg, bg, mask = inverse_mask) composed = cv2.add(masked, bg_masked) target_path = os.path.join(target_dir, 'masked_' + num_str + '.png' ) cv2.imwrite(target_path, composed) import mathlabel_list = ['skin' , 'nose' , 'eye_g' , 'l_eye' , 'r_eye' , 'l_brow' , 'r_brow' , 'l_ear' , 'r_ear' , 'mouth' , 'u_lip' , 'l_lip' , 'hair' , 'hat' , 'ear_r' , 'neck_l' , 'neck' , 'cloth' ] data_root = os.path.join('path_to_project' , 'CelebAMask-HQ' ) example_num = 30000 img_dir = os.path.join(data_root, 'CelebA-HQ-img' ) mask_dir = os.path.join(data_root, 'CelebAMask-HQ-mask-anno' ) total_mask_dir = os.path.join(data_root, 'CelebAMask-total-mask' ) target_dir = os.path.join(data_root, 'CelebAMask-target' ) clean_and_touch_dir(total_mask_dir) clean_and_touch_dir(target_dir) parallel_num = 20 batch_num = math.ceil(example_num / parallel_num) threadings = [] for i in range(parallel_num): start = batch_num * i end = min(batch_num*(i+1 ), example_num) print("pending batch examples from {} to {}" .format(start, end)) t = threading.Thread(target=work_on_imgs, args=(start,end, )) threadings.append(t) t.start() for t in threadings: t.join()
使用脚本转为 tfrecords 或许可以将文件按照 Pascal VOC 的组织方式整理下，然后用 deeplab 的脚本转成训练数据。不过，由于我这里已经有一套自己用得比较熟悉的 tfrecords 生成脚本，改下 file_path 到 feature 的转换函数就行。我的 tfrecords 生成工具脚本在这里：
https://gist.github.com/ec14e54f8b9ee5ce7e4b96e3f39d9302#file-parallel_tfrecords-py
针对本数据集，更换 parser 如下:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 #deeplab 使用 decode_jpeg 或 decode_png operator 解析图像数据，因此生成 feature 的时候，要注意传入 png 编码过的数据。 import io def parse_example_deeplab_style(file_path): filename = os.path.basename(file_path) img_format = filename.split('.')[-1] image = Image.open(file_path) image = image.convert('RGB') image = image.resize((TARGET_SIZE, TARGET_SIZE)) imgByteArr = io.BytesIO() image.save(imgByteArr, format='PNG') image_data = imgByteArr.getvalue() index = os.path.splitext(filename)[0] mask_path = os.path.join(mask_dir, index + '.png') # seg_data = tf.gfile.FastGFile(mask_path, 'rb').read() # the masks are already in (512, 512) size mask_img = Image.open(mask_path) mask_img = mask_img.resize((TARGET_SIZE, TARGET_SIZE)) mask_np = np.uint8(np.array(mask_img)) portrait_mask_np = np.uint8(mask_np>0) # get portrait mask, which only cares about face or background portrait_mask_image = Image.fromarray(portrait_mask_np) # convert back to PIL Image maskByteArr = io.BytesIO() portrait_mask_image.save(maskByteArr, format='PNG') portrait_mask_data = maskByteArr.getvalue() # save to BytesIO and get encoded values without disk writingg tf_example = tf.train.Example(features=tf.train.Features(feature={ 'image/encoded': _bytes_list_feature(image_data), 'image/filename': _bytes_list_feature(filename), 'image/format': _bytes_list_feature(_IMAGE_FORMAT_MAP[img_format]), 'image/height': _int64_list_feature(TARGET_SIZE), 'image/width': _int64_list_feature(TARGET_SIZE), 'image/channels': _int64_list_feature(3), 'image/segmentation/class/encoded': (_bytes_list_feature(portrait_mask_data)), 'image/segmentation/class/format': _bytes_list_feature('png'), })) return tf_example
训练以及验证 首先要针对自己的数据集修改 dataset_generator.py 中的代码，增加数据集参数，主要是 num_classes 及 train/eval 样本个数。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 _CELEBAMASK_FACIAL_SEG_INFORMATION = DatasetDescriptor( splits_to_sizes={ 'train': 28480, 'val': 1498, }, num_classes=2, ignore_label=255, ) _DATASETS_INFORMATION = { 'cityscapes': _CITYSCAPES_INFORMATION, 'pascal_voc_seg': _PASCAL_VOC_SEG_INFORMATION, 'ade20k': _ADE20K_INFORMATION, 'celebamask_facial_seg': _CELEBAMASK_FACIAL_SEG_INFORMATION, }
Train 训练 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 python deeplab/train.py \ --logtostderr \ --training_number_of_steps=300000 \ --train_split="train" \ --model_variant="mobilenet_v2" \ --atrous_rates=6 \ --atrous_rates=12 \ --atrous_rates=18 \ --output_stride=16 \ --decoder_output_stride=4 \ --train_crop_size="513,513" \ --train_batch_size=16 \ --dataset="celebamask_facial_seg" \ --train_logdir=path_to_proj_portrait2anime/train_log_deeplab/celebamask_facial_seg_mobilenetv2_513_bs_16_08221853 \ --dataset_dir=path_to_proj_portrait2anime/CelebAMask-HQ/facial_segmentation_tfrecords_for_deeplab_0818_131910
在我测试下来，对于我这里的人脸分割任务，大概要至少1小时后才能有一定结果。
Eval 验证 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 python deeplab/eval.py \ --logtostderr \ --eval_split="val" \ --model_variant="mobilenet_v2" \ --atrous_rates=12 \ --atrous_rates=24 \ --atrous_rates=36 \ --output_stride=8 \ --decoder_output_stride=4 \ --eval_crop_size="513,513" \ --train_batch_size=8 \ --dataset="celebamask_facial_seg" \ --checkpoint_dir=path_to_proj_portrait2anime/deeplab_train_log/celebamask_facial_seg_mobilenetv2_513_8_08181604 \ --eval_logdir=path_to_proj_portrait2anime/deeplab_train_log/eval_0818_131910 \ --dataset_dir=path_to_proj_portrait2anime/CelebAMask-HQ/facial_segmentation_tfrecords_for_deeplab_0818_131910
Eval 里各参数基本与 train 时一直，但 train_crop_size 要改为 eval_crop_size。另外我在某次实验时，训练的 train_crop_size 调低到 449，验证时 eval_crop_size 还用这个值会报错，调到 513 后正常。这里有解释 ，验证时用全图大小就行，似乎不用顾虑显存/模型大小，只是我暂时还没完全理解。
另外稍提下，Eval 是启动进程不断跟踪 eval_logdir 下的 ckpt 并持续验证，所以不会直接打印结果，需要进入 tensorboard 中看 miou 结果。
Vis 可视化 1 2 3 4 5 6 7 8 9 10 11 12 13 14 python deeplab/vis.py \ --logtostderr \ --vis_split="val" \ --model_variant="mobilenet_v2" \ --atrous_rates=12 \ --atrous_rates=24 \ --atrous_rates=36 \ --output_stride=8 \ --decoder_output_stride=4 \ --vis_crop_size="513,513" \ --dataset="celebamask_facial_seg" \ --checkpoint_dir=path_to_proj_portrait2anime/deeplab_train_log/celebamask_facial_seg_mobilenetv2_513_8_08181604 \ --vis_logdir=path_to_proj_portrait2anime/deeplab_train_log/vis_0818_131910 \ --dataset_dir=path_to_proj_portrait2anime/CelebAMask-HQ/facial_segmentation_tfrecords_for_deeplab_0818_131910
结果类似下图（训练时间短）：
Export 部署 .pb and SavedModel deeplab 项目中也提供了 export_model 代码，用于固化成 .pb 模型。模型固化前后节点信息变化不少，固化前输入输出节点难找，因此这里为了得到 SavedModel 干脆先转 .pb 固化模型，代码里比较明确地给出了输入输出节点名，这样直接拿来用会方便很多：
1 2 _INPUT_NAME = 'ImageTensor' _OUTPUT_NAME = 'SemanticPredictions'
不过这个 export_model.py 的参数与之前有些差距，容易踩坑，先上正确的例子：
1 2 3 4 5 6 7 8 9 10 11 12 python deeplab/export_model.py \ --model_variant="mobilenet_v2" \ --crop_size=513 \ --crop_size=513 \ --atrous_rates=6 \ --atrous_rates=12 \ --atrous_rates=18 \ --output_stride=16 \ --num_classes=2 \ --decoder_output_stride=4 \ --checkpoint_path=path_to_proj_portrait2anime/train_log_deeplab/celebamask_facial_seg_mobilenetv2_513_8_08181604/model.ckpt-15016 \ --export_path=path_to_proj_portrait2anime/train_log_deeplab/frozen/frozen_inference_graph_mobilenetv2_test233.pb
以上面命令来讲，虽然使用了 mobilenet_v2 backbone，但是训练的时候如果设置了 atrous_rates，export时注意别漏掉；之前训练时要指定 --dataset，这里改用 --num_classes 来确定模型输出结构；crop_size 在 export_model.py 里定义为 DEFINE_multi_integer，所以要像这样改用多次列出的格式。
最终效果样例 使用 mobilenet_v2、 output_stride=16、batch_size=16 训练一晚达到 step-88170，loss在约0.4，难以下降，验证 miou 只有 0.83。
loss 过程如下：
实际用在项目中，部分典型结果如下：
跟训练数据集相似的：
边缘清晰，侧脸：
比较难的：
一言难尽，miou 明显不够，实际结果也确实难以达到要求。
