本文主要介绍如何使用自己的数据集训练DeepLabv3+分割算法,代码使用的是官方源码。
1、代码简介
当前使用TensorFlow版本的官方源码,选择它的原因是因为代码中的内容比较全面,除了代码实现以外,还提供了许多文档帮助理解与使用,同时还提供了模型转换的代码实现。
代码地址: 【github】models/research/deeplab at master · tensorflow/models
接下来,先对这个代码仓库进行一下简单的介绍,因为自己在使用该代码仓库的时候只关心训练代码的实现,而忽略的其他的内容,走了不少弯路,到后面才发现我想要的内容,仓库里面早有(==)。
在当前的实现中,我们支持采用以下网络主干:
- MobileNetv2和MobileNetv3:一个为移动设备设计的快速网络结构
- Xception:用于服务器端部署的强大网络结构
- ResNet-v1-{50, 101}:我们提供原始的ResNet-v1及其“ beta”变体,其中对“ stem”进行了修改以进行语义分割。
- PNASNet: 一个通过神经体系结构搜索发现的强大网络结构。
- Auto-Deeplab(代码中叫做HNASNet):通过神经体系结构搜索找到的特定于细分的网络主干。
该目录包含TensorFlow 实现。我们提供的代码使用户可以训练模型,根据mIOU(平均交叉点求和)评估结果以及可视化细分结果。我们以PASCAL VOC 2012和Cityscapes语义分割基准为例。
代码中几个重要文件:
- datasets/:该文件夹下包含对于训练数据集的处理代码,主要针对 PASCAL VOC 2012和Cityscapes数据集的处理。
- g3doc/:该文件夹下包含多个Markdown文件,非常有用,如何安装,常见问题FAQ等。
- deeplab_demo.ipynb:该文件中给出了如果对一张图像进行语义分割并显示结果的Demo。
- export_model.py:该文件提供了将训练的checkpoint模型转为.pb文件的代码实现。
- train.py:训练代码文件,训练时,需要指定提供的训练参数。
- eval.py:验证代码,输出mIOU,用来评估模型的好坏。
- vis.py:可视化代码。
2、安装
Deeplab依赖的库有:
- Numpy
- Pillow 1.0
- tf Slim (which is included in the “tensorflow/models/research/” checkout)
- Jupyter notebook
- Matplotlib
- Tensorflow
2.1 添加库到PYTHONPATH
本地运行的时候,tensorflow/models/research/目录应该追加到PYTHONPATH中,如下:
# From tensorflow/models/research/
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
# [Optional] for panoptic evaluation, you might need panopticapi:
# https://github.com/cocodataset/panopticapi
# Please clone it to a local directory ${PANOPTICAPI_DIR}
touch ${PANOPTICAPI_DIR}/panopticapi/__init__.py
export PYTHONPATH=$PYTHONPATH:${PANOPTICAPI_DIR}/panopticapi
注意:此命令需要在您启动的每个新终端上运行。如果希望避免手动运行此命令,可以将它作为新行添加到〜/ .bashrc文件的末尾。
2.2 测试是否安装成功
通过运行model_test.py快速测试:
# From tensorflow/models/research/
python deeplab/model_test.py
在PASCAL VOC 2012数据集上快速运行所有代码:
# From tensorflow/models/research/deeplab
sh local_test.sh
3、数据集准备
最终目标: 生成TFRecord格式的数据
数据集目录结构如下:
+dataset #数据集名称
+image
+mask
+index
- train.txt
- trainval.txt
- val.txt
+tfrecord
- image: 原图图像,RGB彩色图像
- mask:像素值为类别标签的mask图像,单通道,与原图的名称一致,后缀为.jpg和.png都可以,只要在代码中读取一致即可。VOC数据集默认原图是.jpg,mask图像为.png。
- index:存放图像文件名的txt文件(不加后缀)
- tfrecord:存放转为tfrecord格式的图像数据
数据集制作流程:
- 标注数据,制作符合要求的mask图像
- 将数据集分割为训练集、验证集和测试集
- 生成TFRecord格式的数据集
3.1 标注数据
训练集数据包含两部分,一是原图,二是对应分类的标注值(本文中称为mask图像)。
mask图像的值是如何设置的? 根据图像分割的分类个数来制作原图对应的mask图像。假如一共有N个类别(背景作为一类),则mask图像的值的范围是[0~N)。0值作为背景值,其他分割类别的值依次设置为1, 2, ..., N-1。
注意:
- ignore_label:从字面意思来讲是忽略的标签,即ignore_label是指没有做标注的像素,即不需进行预测的像素值,因此,它不参与loss值的计算,在mask图像中将其值记为255。
- mask图像是单通道的灰度图像。
- mask图像的格式没有限定,但所有的mask图像采用同一种图像格式,方便数据读取。
小总结mask图像的值分为三类:
- 背景:用0表示
- 分类类别:使用1, 2, ....., N-1表示
- ignore_label值:用255表示
如果分割的类别较少,则生成的mask图像看上去是一片黑,因为分类的值都较小,在0~255的范围内不容易显示出来。
3.2 分割数据集
这部分就是将准备的数据集进行分割,分为训练集、验证集、测试集。 无需将具体的图像文件分到三个文件夹中,只需要建立图像的索引文件即可,通过添加相应的路径+文件名即可获取到具体的图像。
假设原图像和mask图像的存放路径如下:
- 原图:./dataset/images
- mask图像:./dataset/mask:此处存放的是2.1小节要求格式
原图与mask图像是一一对应的,包括图像尺寸,图像名(后缀可以不同)
索引文件存放路径:./dataset/index,该路径下生成:
- train.txt
- trainval.txt
- val.txt
索引文件中,只需记录文件名(不加后缀),这取决于代码中数据集加载的方式。
目前为止,数据集目录结构如下:
#./dataset
+image
+mask
+index
- train.txt
- trainval.txt
- val.txt
3.3 将数据打包为TFRecord格式
TFRecord是谷歌推荐的一种二进制文件格式,理论上它可以保存任何格式的信息。TFRecord内部使用了“Protocol Buffer”二进制数据编码方案,它只占用一个内存块,只需要一次性加载一个二进制文件的方式即可,简单,快速,尤其对大型训练数据很友好。而且当我们的训练数据量比较大的时候,可以将数据分成多个TFRecord文件,来提高处理效率。
那么,如何将数据生成TFRecord格式呢?
在此,我们可以借助 项目代码中./datasets/build_voc2012_data.py文件来实现。给文件是VOC2012数据集处理的代码,我们只需修改一下输入参数即可。
参数:
- image_folder:原图文件夹名称,./dataset/image
- semantic_segmentation_folder:分割文件夹名称, ./dataset/mask
- list_folder:索引文件夹名称,./dataset/index
- output_dir:输出路径,即生成的tfrecord文件所在位置,./dataset/tfrecord
运行命令:
python ./datasets/build_voc2012_data.py --image_folder=./dataset/image
--semantic_segmentation_folder=./dataset/mask
--list_folder=./dataset/index
--output_dir=./dataset/tfrecord
生成的文件如下:
注意: 可在代码中调节参数_NUM_SHARDS (默认为4),改变数据分块的数目。(一些文件系统有最大单个文件大小的限制,如果数据集非常大,增加_NUM_SHARDS 可减小单个文件的大小)
该文件的核心代码如下:
# dataset_split指的是train.txt, val.txt等
dataset = os.path.basename(dataset_split)[:-4]
filenames = [x.strip('\n') for x in open(dataset_split, 'r')] # 文件名列表
# 输出tfrecord文件名
output_filename = os.path.join(
FLAGS.output_dir,
'%s-%05d-of-%05d.tfrecord' % (dataset, shard_id, _NUM_SHARDS))
with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:
for i in range(start_idx, end_idx):
image_filename = os.path.join(iamge_folder, filenames[i]+'.'+image_format)# 原图路径
image_data = tf.gfile.GFile(image_filename, 'rb').read() #读取原图文件
height, width = image_reader.read_image_dims(image_data)
seg_filename = os.path.join(semantic_segmentation_folder,
filenames[i] + '.' + label_format) # mask图像路径
seg_data = tf.gfile.GFile(seg_filename, 'rb').read() # 读取分割图像
seg_height, seg_width = label_reader.read_image_dims(seg_data)
# 判断原图与mask图像尺寸是否匹配
if height != seg_height or width != seg_width:
raise RuntimeError('Shape mismatched between image and label.')
# Convert to tf example.
example = build_data.image_seg_to_tfexample(
image_data, filenames[i], height, width, seg_data)
tfrecord_writer.write(example.SerializeToString())
至此,数据集的制作部分已经完成!!!
4、训练
4.1 代码修改
为了训练自己的数据集,需要修改以下几处文件:
1 datasets/data_generator.py:增加数据集的注册
该文件提供语义分割数据的包装器
在该文件中,可以看到PASCAL_VOC, CITYSCAPES以及ADE20K数据集的数据描述,如下:
_PASCAL_VOC_SEG_INFORMATION = DatasetDescriptor(
splits_to_sizes={
'train': 1464,
'train_aug': 10582,
'trainval': 2913,
'val': 1449,
},
num_classes=21,
ignore_label=255,
)
en,比着葫芦画瓢,增加我们自己数据集的描述信息,如下:
_PORTRAIT_INFORMATION = DatasetDescriptor(
splits_to_sizes={
'train': 17116,
'trainval': 21395,
'val': 4279,
},
num_classes=2, # 类别数目,包括背景
ignore_label=255, # 忽略像素值
)
以人像分割任务为例,只有两类,即前景(人像)和背景(非人像)。
添加完描述信息后,需要将该数据集信息进行注册,如下:
_DATASETS_INFORMATION = {
'cityscapes': _CITYSCAPES_INFORMATION,
'pascal_voc_seg': _PASCAL_VOC_SEG_INFORMATION,
'ade20k': _ADE20K_INFORMATION,
'portrait_seg': _PORTRAIT_INFORMATION, #增加此句
}
注意:此处的数据集名称要与前面对应!
2 ./utils/train_utils.py修改
在函数get_model_init_fn中,修改为如下代码,增加logits层不加载预训练模型权重:
# Variables that will not be restored.
exclude_list = ['global_step', 'logits']
if not initialize_last_layer:
exclude_list.extend(last_layers)
4.2 主要训练参数
训练文件train.py和common.py文件中包含了训练分割网络所需要的所有参数。
- model_variant:Deeplab模型变量,可选值可见core/feature_extractor.py。
- 当使用mobilenet_v2时,设置变量strous_rates=decoder_output_stride=None;
- 当使用xception_65或resnet_v1时,设置strous_rates=[6,12,18](output stride 16), decoder_output_stride=4。
- label_weights:此变量可以设置标签的权重值,当数据集中出现类别不均衡时,可通过此变量来指定每个类别标签的权重值,如label_weights=[0.1, 0.5]意味着标签0的权重是0.1, 标签1的权重是0.5。如果该值为None,则所有的标签具有相同的权重1.0。
- train_logdir:存放checkpoint和logs的路径。
- log_steps:该值表示每隔多少步输出日志信息。
- save_interval_secs:该值表示以秒为单位,每隔多长时间保存一次模型文件到硬盘。
- optimizer:优化器,可选值['momentum', 'adam']。
- learning_policy:学习率策略,可选值['poly', 'step']。
- base_learning_rate:基础学习率,默认值0.0001。
- training_number_of_steps:模型训练的迭代次数。
- train_batch_size:模型训练的批处理图像数量。
- train_crop_size:模型训练时所使用的图像尺寸,默认'513, 513'。
- tf_initial_checkpoint:预训练模型。
- initialize_last_layer:是否初始化最后一层。
- last_layers_contain_logits_only:是否只考虑逻辑层作为最后一层。
- fine_tune_batch_norm:是否微调batch norm参数。
- atrous_rates:默认值[6, 12, 18]。
- output_stride:默认值16,输入和输出空间分辨率的比值
- 对于xception_65, 如果output_stride=8,则使用atrous_rates=[12, 24, 36]
- 如果output_stride=16,则atrous_rates=[6, 12, 18]
- 对于mobilenet_v2,使用None
- 注意:在训练和验证阶段可以使用不同的strous_rates和output_stride。
- dataset:所使用的分割数据集,此处与数据集注册时的名称一致。
- train_split:使用哪个数据集来训练,可选值即数据集注册时的值,如train, trainval。
- dataset_dir:数据集存放的路径。
针对训练参数,下面几点需要重点注意:
- 关于是否加载预训练网络的权重问题 如果要在其他数据集上微调该网络,需要关注以下几个参数:
- 使用预训练网络的权重,设置initialize_last_layer=True
- 只使用网络的backbone,设置initialize_last_layer=False和last_layers_contain_logits_only=False
- 使用所有的预训练权重,除了logits,设置initialize_last_layer=False和last_layers_contain_logits_only=True
由于我的数据集分类与默认类别数不同,因此采取的参数值是:
--initialize_last_layer=false
--last_layers_contain_logits_only=true
- 如果资源有限,想要训练自己数据集的几条建议:
- 设置output_stride=16或者甚至32(同时需要修改atrous_rates变量,例如,对于output_stride=32,atrous_rates=[3, 6, 9])
- 尽可能多的使用GPU,更改num_clone标志,并将train_batch_size设置的尽可能大
- 调整train_crop_size,可以将它设置的更小一些,例如513x513(甚至321x321),这样就可以使用更大的batch_size
- 使用较小的网络主干,如mobilenet_v2
- 关于是否微调batch_norm 当训练使用的批处理大小train_batch_size大于12(最好大于16)时,设置fine_tune_batch_norm=True。否则,设置fine_tune_batch_norm=False。
4.3 预训练模型
模型链接具体可见:models/model_zoo.md at master · tensorflow/models
提供了在几个数据集上的预训练模型,包括(1) PASCAL VOC 2012, (2) Cityscapes, (3) ADE20K
未解压的目下包括:
- 一个frozen inference graph(forzen_inference_graph.pb)。默认情况下,所有冻结推理图的输出步长为8,单个eval scale为1.0,没有左右翻转,除非另外指定。基于MobileNet-v2的模型不包括解码器模块。
- 一个checkpoint(model.ckpt.data-00000-of-00001, model.ckpt.index)
还提供了在ImageNet预训练的checkpoints
未解压文件包括:
一个model checkpoint (model.ckpt.data-00000-of-00001, model.ckpt.index)
根据自己的情况进行下载
4.4 训练模型
python train.py \
--logtostderr \
--training_number_of_steps=20000 \
--train_split="train" \
--model_variant="xception_65" \
--train_crop_size="513,513" \
--atrous_rates=6 \
--atrous_rates=12 \
--atrous_rates=18 \
--output_stride=16 \
--decoder_output_stride=4 \
--train_batch_size=2 \
--save_interval_secs=240 \
--optimizer="momentum" \
--leraning_policy="poly" \
--fine_tune_batch_norm=false \
--initialize_last_layer=false \
--last_layers_contain_logits_only=true \
--dataset="portrait_seg" \
--tf_initial_checkpoint="./checkpoint/deeplabv3_pascal_trainval/model.ckpt" \
--train_logdir="./train_logs" \
--dataset_dir="./dataset/tfrecord"
4.5 验证模型
验证代码: ./eval.py
# From tensorflow/models/research/
python deeplab/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="513,513" \
--dataset="portrait_seg" \ # 数据集名称
--checkpoint_dir=${PATH_TO_CHECKPOINT} \ # 预训练模型
--eval_logdir=${PATH_TO_EVAL_DIR} \
--dataset_dir="./dataset/tfrecord" # 数据集路径
得到的结果如下:
4.6 训练过程可视化
可以使用Tensorboard检查培训和评估工作的进展。如果使用推荐的目录结构,Tensorboard可以使用以下命令运行:
tensorboard --logdir=${PATH_TO_LOG_DIRECTORY}
# 文中log地址
tensorboard --logdir="./train_logs"
5、推理
5.1 模型导出
在训练过程中,会保存模型文件到硬盘,如下:
其形式是TensorFlow的checkpoint格式,代码中提供了一个脚本(export_model.py)可以将checkpoint转换为.pb格式。
export_model.py主要参数:
- checkpoint_path:训练保存的检查点文件
- export_path:模型导出路径
- num_classes:分类类别
- crop_size:图像尺寸,[513, 513]
- atrous_rates:12, 24, 36
- output_stride:8
生成的.pb文件如下:
5.2 单张图像上推理
class DeepLabModel(object):
"""class to load deeplab model and run inference"""
INPUT_TENSOR_NAME = 'ImageTensor:0'
OUTPUT_TENSOR_NAME='SemanticPredictions:0'
INPUT_SIZE = 513
FROZEN_GRAPH_NAME= 'frozen_inference_graph'
def __init__(self, pretrained_weights):
"""Creates and loads pretrained deeplab model."""
self.graph = tf.Graph()
graph_def = None
# Extract frozen graph from tar archive
if pretrained_weights.endswith('.tar.gz'):
tar_file = tarfile.open(pretrained_weights)
for tar_info in tar_file.getmembers():
if self.FROZEN_GRAPH_NAME in os.path.basename(tar_info.name):
file_handle = tar_file.extractfile(tar_info)
graph_def = tf.GraphDef.FromString(file_handle.read())
break
tar_file.close()
else:
with open(pretrained_weights, 'rb') as fd:
graph_def = tf.GraphDef.FromString(fd.read())
if graph_def is None:
raise RuntimeError('Cannot find inference graph in tar archive.')
with self.graph.as_default():
tf.import_graph_def(graph_def, name='')
gpu_options = tf.GPUOptions(allow_growth=True)
config = tf.ConfigProto(gpu_options=gpu_options, log_device_placement=False)
self.sess = tf.Session(graph=self.graph, config=config)
def run(self, image):
"""Runs inference on a single image.
Args:
image: A PIL.Image object, raw input image.
Returns:
resized_image:RGB image resized from original input image.
seg_map:Segmentation map of 'resized_iamge'.
"""
width, height = image.size
resize_ratio = 1.0 * self.INPUT_SIZE/max(width, height)
target_size = (int(resize_ratio*width), int(resize_ratio * height))
resized_image = image.convert('RGB').resize(target_size, Image.ANTIALIAS)
batch_seg_map = self.sess.run(
self.OUTPUT_TENSOR_NAME,
feed_dict={self.INPUT_TENSOR_NAME:[np.asarray(resized_image)]}
)
seg_map = batch_seg_map[0]
return resized_image, seg_map
if __name__ == '__main__':
pretrained_weights = './train_logs/frozen_inference_graph_20000.pb'
MODEL = DeepLabModel(pretrained_weights) # 加载模型
img_name = 'test.jpg'
img = Image.open(img_name)
resized_im, seg_map = MODEL.run(img) #获取结果
seg_map[seg_map==1]=255 #将人像的像素值置为255
seg_map= Image.fromarray(seg_map.astype('uint8'))
seg_map.save('output.jpg') # 保存mask结果图像
至此,整个训练过程就结束了!!!