DeepLabV3+模型训练—学习笔记

环境准备

添加依赖库

# From tensorflow/models/research/
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
export CUDA_VISIBLE_DEVICES=1

测试环境

# From tensorflow/models/research/
python3 deeplab/model_test.py

数据准备

数据处理

# !/usr/bin/python
# -*- coding = utf8 -*-
"""
# @ Author: ***
# @ Created Time: 2019-05-13 09:31:37
# @ Description: Prepare data for deeplab.
"""

import sys
import os
import argparse
import shutil
import cv2
import numpy as np

def parse_arguments(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_matting',type=str,help='Directory with input \
            images.')
    parser.add_argument('--input_image', type=str, help='Directory with input \
                images.')
    parser.add_argument('--output_dir',type=str,help='Directory with output \
            alpha images.')
    return parser.parse_args(argv)

def main(args):
    output_dir = os.path.expanduser(args.output_dir)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    output_img = os.path.join(output_dir, 'JPEGImages')
    if not os.path.exists(output_img):
        os.makedirs(output_img)
    output_mask = os.path.join(output_dir, 'mask')
    if not os.path.exists(output_mask):
        os.makedirs(output_mask)
    output_seg = os.path.join(output_dir, 'SegmentationClass')
    if not os.path.exists(output_seg):
        os.makedirs(output_seg)
    input_matting = os.path.expanduser(args.input_matting)
    img_dir = os.path.expanduser(args.input_image)

    trainval_filename = os.path.join(output_dir, 'ImageSets', 'Segmentation', 'trainval.txt')
    train_filename = os.path.join(output_dir, 'ImageSets', 'Segmentation', 'train.txt')
    val_filename = os.path.join(output_dir, 'ImageSets', 'Segmentation','val.txt')
    with open(trainval_filename, 'w') as trainval_f, \
            open(train_filename, 'w') as train_f, \
            open(val_filename, 'w') as val_f:
        index = 0
        if os.path.isdir(input_matting):
            for sub1_name in os.listdir(input_matting):
                sub1_dir = os.path.join(input_matting, sub1_name)
                img_sub1_dir = os.path.join(img_dir, sub1_name)

                if not os.path.isdir(sub1_dir):
                    continue
                for sub2_name in os.listdir(sub1_dir):
                    sub2_dir = os.path.join(sub1_dir, sub2_name)
                    img_sub2_dir = os.path.join(img_sub1_dir, sub2_name.replace('matting','clip'))

                    if not os.path.isdir(sub2_dir):
                        continue
                    for image_name in os.listdir(sub2_dir):
                        image_path =  os.path.join(os.path.join(sub2_dir, image_name))
                        orig_image_path = os.path.join(img_sub2_dir, image_name)
                        if not os.path.isfile(orig_image_path):
                            orig_image_path = orig_image_path.replace('png','jpg')
                        print('processing', image_path)

                        image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
                        alpha = image[:,:,3]
                        ret,mask = cv2.threshold(alpha,10,255,cv2.THRESH_BINARY)
                        cv2.imwrite(os.path.join(output_mask, image_name), mask)
                        label_img = np.where(mask == 255, 1, 0)
                        cv2.imwrite(os.path.join(output_seg, image_name), label_img)
                        shutil.copy(orig_image_path, os.path.join(output_img, image_name))

                        index += 1
                        trainval_f.write('%s\n' % image_name.split('.')[0])
                        if index <= 30000:
                            train_f.write('%s\n' % image_name.split('.')[0])
                        else:
                            val_f.write('%s\n' % image_name.split('.')[0])
            


if __name__=='__main__':
    main(parse_arguments(sys.argv[1:]))

数据目录结构

 + portrait  # 自己的数据集名称
    + JPEGImages #RGB original image（.jpg）
    + SegmentationClass #label image（.png）
    + ImageSets
        + Segmentation
            - train.txt # filename without extension
            - val.txt
            - trainval.txt
 + tfrecord

对于SegmentationClass中的lable图像：

ignore_label 如果你的数据集存在ignore_label，注意不要把ignore_label和background混淆。 ignore_label是没有做标注的，不在预测范围内的,ignore_label是不参与计算loss的。在mask中将ignore_label的灰度值标记为255。
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-799qj1Ng-1588058618893)(http://www.pianshen.com/images/3/f293780fc45ea58d1e5193653ac76d2b.png)]

数据不平衡问题
trian_util.py 修改

转换为TFRecord格式

python3 build_voc2012_data.py --image_folder=./portrait/JPEGImages \
                             --semantic_segmentation_folder=./portrait/SegmentationClass \
                             --list_folder=./portrait/ImageSets/Segmentation \ 
                             --output_dir=./portrait/tfrecord

也可以写成shell脚本：

#deeplab/gen_tf.sh
#!/bin/bash
# Update PYTHONPATH.
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim

# Set up the working environment.
CURRENT_DIR=$(pwd)
WORK_DIR="${CURRENT_DIR}/deeplab/datasets"
DATA_DIR="/xxx/data/matting_human_half/portrait"

export CUDA_VISIBLE_DEVICES=1

python3 "${WORK_DIR}"/build_voc2012_data.py \
    --image_folder="${DATA_DIR}"/JPEGImages \
    --semantic_segmentation_folder="${DATA_DIR}/SegmentationClass" \
    --list_folder="${DATA_DIR}/ImageSets/Segmentation" \
    --output_dir="${DATA_DIR}/tfrecord"

训练

注册数据集

之前的版本是修改segmentation_dataset.py

#deeplab/datasets/data_generator.py
_PORTRAIT_INFORMATION = DatasetDescriptor(
    splits_to_sizes={
        'train':30000, 
        'val':4426,
        'trainval':34426
    },
    num_classes=2, #classes include background
    ignore_label=255,
)

_DATASETS_INFORMATION = {
    'cityscapes': _CITYSCAPES_INFORMATION,
    'pascal_voc_seg': _PASCAL_VOC_SEG_INFORMATION,
    'ade20k': _ADE20K_INFORMATION,
    'portrait': _PORTRAIT_INFORMATION,
}

训练参数设置

logits层是在softmax前的最后一层特征图，是没有经过归一化的预测值，如果对分类问题再经过一层softmax就可以得到每类的概率。
或者直接修改代码，在使用预训练权重的时候，不加载该logits层。

#deeplab/utils/train_utils.py
# Variables that will not be restored.
#exclude_list = ['global_step']
exclude_list = ['global_step', 'logits']

# When fine_tune_batch_norm=True, use at least batch size larger than 12
# (batch size more than 16 is better). Otherwise, one could use smaller batch
# size and set fine_tune_batch_norm=False.

# For `xception_65`, use atrous_rates = [12, 24, 36] if output_stride = 8, or
# rates = [6, 12, 18] if output_stride = 16. For `mobilenet_v2`, use None. Note
# one could use different atrous_rates/output_stride during training/evaluation.

数据不平衡问题

train_utils.py中已无修改代码。

训练

# deeplab/train_portrait.sh
#!/bin/bash
# Update PYTHONPATH.
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim

# Set up the working environment.
CURRENT_DIR=$(pwd)
WORK_DIR="${CURRENT_DIR}/deeplab"
DATA_DIR="/mnt/data/matting_human_half/portrait"

export CUDA_VISIBLE_DEVICES=1

EXP_FOLDER="exp/train_on_trainval_set"
TRAIN_LOGDIR="${DATA_DIR}/${EXP_FOLDER}/train"
EVAL_LOGDIR="${DATA_DIR}/${EXP_FOLDER}/eval"
VIS_LOGDIR="$${DATA_DIR}/${EXP_FOLDER}/vis"
EXPORT_DIR="$${DATA_DIR}/${EXP_FOLDER}/export"
mkdir -p "${TRAIN_LOGDIR}"
mkdir -p "${EVAL_LOGDIR}"
mkdir -p "${VIS_LOGDIR}"
mkdir -p "${EXPORT_DIR}"

python3 "${WORK_DIR}"/train.py \
 --logtostderr \
 --train_split="train" \
 --model_variant="xception_65" \
 --atrous_rates=6 \
 --atrous_rates=12 \
 --atrous_rates=18 \
 --output_stride=16 \
 --decoder_output_stride=4 \
 --train_crop_size=513 \
 --train_crop_size=513 \
 --train_batch_size=2 \
 --dataset="portrait" \
 --training_number_of_steps=30000 \
 --initialize_last_layer=False \
 --last_layers_contain_logits_only=True \
 --fine_tune_batch_norm=False \
 --tf_initial_checkpoint="${WORK_DIR}/datasets/pascal_voc_seg/init_models/deeplabv3_pascal_train_aug/model.ckpt" \
 --train_logdir="${TRAIN_LOGDIR}" \
 --dataset_dir="${DATA_DIR}/tfrecord"

训练过程可视化

在tf.summary里设置好要查看保存的参数后运行会生成events.out.tfevents.{time}.{machine-name}的文件，这个就是用tensorboard来查看的日志文件。当训练和验证的日志文件不在同一个文件中时，想同时查看两个日志，要转到上一层目录，使用下面的命令，在浏览器中访问 http://服务器名：6006 即可查看。

tensorboard --logdir /xxx/data/matting_human_half/portrait/exp/train_on_trainval_set

验证

下载的版本运行eval.py不会出现mIoU分数，在summary里面加入tf.Print的op就可以显示了。

print_miou = tf.Print(miou, [miou], 'miou is :') #create print op
tf.summary.scalar('print_miou', print_miou)

summary_op = tf.summary.merge_all()

 python3 "${WORK_DIR}"/eval.py \
  --logtostderr \
  --eval_split="val" \
  --model_variant="xception_65" \
  --atrous_rates=6 \
  --atrous_rates=12 \
  --atrous_rates=18 \
  --output_stride=16 \
  --decoder_output_stride=4 \
  --eval_crop_size=801 \
  --eval_crop_size=601 \
  --dataset="portrait" \
  --checkpoint_dir="${TRAIN_LOGDIR}" \
  --eval_logdir="${EVAL_LOGDIR}" \
  --dataset_dir="${DATA_DIR}/tfrecord" \
  --max_number_of_evaluations=1

注意：–eval_crop_size必须是输入测试图像中最大的图像尺寸，否则会出错。如果测试集中图像的最大宽或最大高为800，那么–eval_crop_size=801。

可视化

使用版本中已无需修改，只能从下面两个中选择。

flags.DEFINE_enum('colormap_type', 'pascal', ['pascal', 'cityscapes'],
                  'Visualization colormap type.')

python3 "${WORK_DIR}"/vis.py \
 --logtostderr \
 --vis_split="val" \
 --model_variant="xception_65" \
 --atrous_rates=6 \
 --atrous_rates=12 \
 --atrous_rates=18 \
 --output_stride=16 \
 --decoder_output_stride=4 \
 --vis_crop_size=801 \
 --vis_crop_size=601 \
 --dataset="portrait" \
 --checkpoint_dir="${TRAIN_LOGDIR}" \
 --vis_logdir="${VIS_LOGDIR}" \
 --dataset_dir="${DATA_DIR}/tfrecord" \
 --max_number_of_iterations=1

注意：–vis_crop_size同–eval_crop_size。

模型导出

# CKPT_PATH="${TRAIN_LOGDIR}/model.ckpt-${NUM_ITERATIONS}"
# EXPORT_PATH="${EXPORT_DIR}/frozen_inference_graph.pb"

# python3 "${WORK_DIR}"/export_model.py \
  # --logtostderr \
  # --checkpoint_path="${CKPT_PATH}" \
  # --export_path="${EXPORT_PATH}" \
  # --model_variant="xception_65" \
  # --atrous_rates=6 \
  # --atrous_rates=12 \
  # --atrous_rates=18 \
  # --output_stride=16 \
  # --decoder_output_stride=4 \
  # --num_classes=2 \
  # --crop_size=801 \
  # --crop_size=601 \
  # --inference_scales=1.0

测试

测试输出的数据类型为（h,w,1），数据类型为int64，需要进一步处理需要类型转换。

import sys
import os
import time
import argparse
import tensorflow as tf
import numpy as np
import cv2


def parse_arguments(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('--image_dir',type=str,help='input images  dir', \
                         default='./data/test_ids')
    parser.add_argument('--model_dir',type=str,help='input model  dir', \
                        default='./portrait/exp/train_on_trainval_set_0007/export/frozen_inference_graph.pb')
    parser.add_argument('--output_dir',type=str,help='output images  dir', \
                        default='./portrait/output')

    return parser.parse_args(argv)
def prepare_image(img):
    h, w, _ = img.shape
    print(img.shape)
    if h >= w:
        ratio = 800. / h
        img_tmp = cv2.resize(img, None, fx=ratio, fy=ratio, interpolation=cv2.INTER_AREA)
        h_new, w_new, _ = img_tmp.shape
        print(img_tmp.shape)
        if w_new > 600:
            ratio_new = 600. / w_new
            img_resized = cv2.resize(img_tmp, None, fx=ratio_new, fy=ratio_new, interpolation=cv2.INTER_AREA)
        else:
            img_resized = img_tmp
    else:
        ratio = 600. / w
        img_resized = cv2.resize(img, None, fx=ratio, fy=ratio, interpolation=cv2.INTER_AREA)

    return img_resized

def segment(net_fun, image):
    img = prepare_image(image)
    img_input = np.expand_dims(img, axis=0).astype(np.uint8)
    result = net_fun(img_input)
    print(result.shape)
    mask = np.where(result.transpose((1, 2, 0)) == 1, 255, 0)
    seg_img = np.c_[img, mask]
    return seg_img

def main(args):
    output_dir = os.path.expanduser(args.output_dir)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    graph = tf.Graph()
    graph_def = None
    with open(args.model_dir, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    if graph_def is None:
        raise RuntimeError('Cannot find inference graph in tar archive.')
    with graph.as_default():
        tf.import_graph_def(graph_def, name='')
        sess = tf.Session(graph=graph, config=tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True)))
        input_tensor = graph.get_tensor_by_name('ImageTensor:0')
        output_tensor = graph.get_tensor_by_name('SemanticPredictions:0')
        with sess.as_default():
            net_fun = lambda img : sess.run(output_tensor, feed_dict={input_tensor:img})

    if os.path.isfile(args.image_dir):
        image = cv2.imread(args.image_dir)
        image_name = os.path.basename(args.image_dir).split('.')[0] + '.png'
        output_image = os.path.join(args.output_dir, image_name)
        seg_img = segment(net_fun, image)
        cv2.imwrite(output_image, seg_img)
    elif os.path.isdir(args.image_dir):
        for image_name in os.listdir(args.image_dir):
            image_path = os.path.join(args.image_dir, image_name)
            output_image = os.path.join(args.output_dir, image_name.split('.')[0] + '.png')
            image = cv2.imread(image_path)
            seg_img = segment(net_fun, image)
            cv2.imwrite(output_image, seg_img)

if __name__ == '__main__':
    main(parse_arguments(sys.argv[1:]))

源码理解

train.py

• 数据结构初始化
  train.main()函数中
  
  初始化结果
  
• 数据读取
  train.main()函数中
  
  data_generator.get_one_shot_iterator()函数中
  
• TFRecord数据解析
  data_generator._parse_function()函数中
  
  
  
  
  
• 图像预处理
  data_generator._preprocess_image()函数中
  
  input_preprocess.preprocess_image_and_label()
  
• 随机缩放
  
  preprocess_utils.randomly_scale_image_and_lable()
  
• 边界填充
  input_preprocess.preprocess_image_and_label()
  
  feature_extractor.mean_pixel()
  
  preprocess_utils.pad_to_bounding_box()
  
  
• 随机裁剪
  input_preprocess.preprocess_image_and_label()
  
  preprocess_utils.random_crop()
  
  
  preprocess_utils._crop()
  
• 随机翻转
  input_preprocess.preprocess_image_and_label()
  
  preprocess_utils.flip_dim()
  
• 学习率设置
  train._train_deeplab_model()
  
  train_utils.get_model_learning_rate()
  
  
• 损失函数计算
  train._train_deeplab_model()
  
  train._build_deeplab()
  
  common.new()
  
  
  
  
  train._build_deeplab()
  
  model.multi_scale_logits()
  
  
  model._get_logits()
  
  model.extract_features()
  
  feature_extractor.extract_features()
  
  feature_extractor.arg_scopes_map()
  
  xception.xception_arg_scope()
  
  
  feature_extractor.extract_features()
  
  feature_extractor.get_network()
  
  feature_extractor.networks_map
  
  xception.xception_65()
  
  
  xception.xception_block()
  
  
  
  xception.xception()
  
  resnet_utils.conv2d_same()
  
  arg_scope.add_arg_scope()
  
  xception.xception()
  
  
  xception.stack_blocks_dense()
  
  
  arg_scope.add_arg_scope()
  
  xception.xception_module()
  
  
  xception.xception_module的子函数
  
  
  xception.xception_module()
  
  
  xception.xception()
  
  model.extract_features()
  
  
  
  utils.split_separable_conv2d()
  
  model.extract_features()
  
  model._get_logits()
  
  model.refine_by_decoder()
  
  
  
  
  utils.split_separable_conv2d()
  
  model._get_logits()
  
  
  model.get_branch_logits()
  
  
  model.multi_scale_logits()
  
  train._build_deeplab()
  
  train_utils.add_softmax_cross_entropy_loss_for_each_scale()
  
  
  train._build_deeplab()
  
  train._log_summaries()
  
  train._tower_loss()
  
  train._train_deeplab_model()
  
  train._average_gradients()
  
  train._train_deeplab_model()
  
  model.get_extra_layer_scopes()
  
  train._train_deeplab_model()
  
  train_utils.get_model_gradient_multipliers()
  
  train._train_deeplab_model()
  
  train.mian()
  
  train_utils.get_model_init_fn()
  
  train.mian()
  

参考资料

1. DeepLab V3+ 训练自己的数据
2. 使用 deeplabv3+ 训练自己的数据集经验总结
3. deeplabv3+训练自己的数据集
4. Deeplab V3+训练自己数据集全过程
5. TensorFlow实战：Chapter-9下(DeepLabv3+在自己的数据集训练)
6. TENSORBOARD 可视化网络运行过程
7. DEEPLABV3+ 命令行不显示MIOU的解决