本文主要介绍如何使用自己的数据集训练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.txttrainval.txtval.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/imagesemantic_segmentation_folder：分割文件夹名称,./dataset/masklist_folder：索引文件夹名称，./dataset/indexoutput_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, 36output_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结果图像
至此，整个训练过程就结束了！！！
扫一扫
459
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
