TensorFlow-slim 训练 CNN 分类模型（续）
2018-04-20 本文已影响733人
公输睚信
在前面的文章 TensorFlow-slim 训练 CNN 分类模型 我们已经使用过 tf.contrib.slim 模块来构建和训练模型了，今天我们继续这一话题，但稍有不同的是我们不再定义数据输入输出的占位符，而是使用 tf.contrib.slim 模块来导入 tfrecord 数据进行训练。这样的训练方式主要有两个优点：1.数据是并行读取的，而且并非一次性全部导入到内存，因此可以缓解内存不足的问题；2.使用 tf.contrib.slim 模块的封装函数 slim.learning.train 来训练，可以直接使用 Tensroboard 来监督损失以及准确率等曲线，还可以在中断训练后继续从上次保存的位置继续进行训练。
我们考虑的问题仍然是 10 分类由 Captcha 生成的数字，具体请访问：TensorFlow 训练 CNN 分类器，使用的模型与文章 TensorFlow-slim 训练 CNN 分类模型 的相同。为了方便阅读，再次将模型的代码粘贴如下（命名为：model.py）：
# -*- coding: utf-8 -*-
"""
Created on Fri Mar 30 16:54:02 2018
@author: shirhe-lyh
"""
import tensorflow as tf
from abc import ABCMeta
from abc import abstractmethod
slim = tf.contrib.slim
class BaseModel(object):
"""Abstract base class for any model."""
__metaclass__ = ABCMeta
def __init__(self, num_classes):
"""Constructor.
Args:
num_classes: Number of classes.
"""
self._num_classes = num_classes
@property
def num_classes(self):
return self._num_classes
@abstractmethod
def preprocess(self, inputs):
"""Input preprocessing. To be override by implementations.
Args:
inputs: A float32 tensor with shape [batch_size, height, width,
num_channels] representing a batch of images.
Returns:
preprocessed_inputs: A float32 tensor with shape [batch_size,
height, widht, num_channels] representing a batch of images.
"""
pass
@abstractmethod
def predict(self, preprocessed_inputs):
"""Predict prediction tensors from inputs tensor.
Outputs of this function can be passed to loss or postprocess functions.
Args:
preprocessed_inputs: A float32 tensor with shape [batch_size,
height, width, num_channels] representing a batch of images.
Returns:
prediction_dict: A dictionary holding prediction tensors to be
passed to the Loss or Postprocess functions.
"""
pass
@abstractmethod
def postprocess(self, prediction_dict, **params):
"""Convert predicted output tensors to final forms.
Args:
prediction_dict: A dictionary holding prediction tensors.
**params: Additional keyword arguments for specific implementations
of specified models.
Returns:
A dictionary containing the postprocessed results.
"""
pass
@abstractmethod
def loss(self, prediction_dict, groundtruth_lists):
"""Compute scalar loss tensors with respect to provided groundtruth.
Args:
prediction_dict: A dictionary holding prediction tensors.
groundtruth_lists: A list of tensors holding groundtruth
information, with one entry for each image in the batch.
Returns:
A dictionary mapping strings (loss names) to scalar tensors
representing loss values.
"""
pass
class Model(BaseModel):
"""A simple 10-classification CNN model definition."""
def __init__(self,
is_training,
num_classes):
"""Constructor.
Args:
is_training: A boolean indicating whether the training version of
computation graph should be constructed.
num_classes: Number of classes.
"""
super(Model, self).__init__(num_classes=num_classes)
self._is_training = is_training
def preprocess(self, inputs):
"""Predict prediction tensors from inputs tensor.
Outputs of this function can be passed to loss or postprocess functions.
Args:
preprocessed_inputs: A float32 tensor with shape [batch_size,
height, width, num_channels] representing a batch of images.
Returns:
prediction_dict: A dictionary holding prediction tensors to be
passed to the Loss or Postprocess functions.
"""
preprocessed_inputs = tf.to_float(inputs)
preprocessed_inputs = tf.subtract(preprocessed_inputs, 128.0)
preprocessed_inputs = tf.div(preprocessed_inputs, 128.0)
return preprocessed_inputs
def predict(self, preprocessed_inputs):
"""Predict prediction tensors from inputs tensor.
Outputs of this function can be passed to loss or postprocess functions.
Args:
preprocessed_inputs: A float32 tensor with shape [batch_size,
height, width, num_channels] representing a batch of images.
Returns:
prediction_dict: A dictionary holding prediction tensors to be
passed to the Loss or Postprocess functions.
"""
with slim.arg_scope([slim.conv2d, slim.fully_connected],
activation_fn=tf.nn.relu):
net = preprocessed_inputs
net = slim.repeat(net, 2, slim.conv2d, 32, [3, 3], scope='conv1')
net = slim.max_pool2d(net, [2, 2], scope='pool1')
net = slim.repeat(net, 2, slim.conv2d, 64, [3, 3], scope='conv2')
net = slim.max_pool2d(net, [2, 2], scope='pool2')
net = slim.repeat(net, 2, slim.conv2d, 128, [3, 3], scope='conv3')
net = slim.flatten(net, scope='flatten')
net = slim.dropout(net, keep_prob=0.5,
is_training=self._is_training)
net = slim.fully_connected(net, 512, scope='fc1')
net = slim.fully_connected(net, 512, scope='fc2')
net = slim.fully_connected(net, self.num_classes,
activation_fn=None, scope='fc3')
prediction_dict = {'logits': net}
return prediction_dict
def postprocess(self, prediction_dict):
"""Convert predicted output tensors to final forms.
Args:
prediction_dict: A dictionary holding prediction tensors.
**params: Additional keyword arguments for specific implementations
of specified models.
Returns:
A dictionary containing the postprocessed results.
"""
logits = prediction_dict['logits']
logits = tf.nn.softmax(logits)
classes = tf.cast(tf.argmax(logits, axis=1), dtype=tf.int32)
postprecessed_dict = {'classes': classes}
return postprecessed_dict
def loss(self, prediction_dict, groundtruth_lists):
"""Compute scalar loss tensors with respect to provided groundtruth.
Args:
prediction_dict: A dictionary holding prediction tensors.
groundtruth_lists: A list of tensors holding groundtruth
information, with one entry for each image in the batch.
Returns:
A dictionary mapping strings (loss names) to scalar tensors
representing loss values.
"""
logits = prediction_dict['logits']
loss = tf.reduce_mean(
tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=logits, labels=groundtruth_lists))
loss_dict = {'loss': loss}
return loss_dict
下面重点说明怎么使用 tf.contrib.slim 模块来训练。本文所有代码见 github: slim_cnn_test。
一、slim.learning.train 训练 CNN 模型
我们已经知道通过定义数据入口、出口的占位符 tf.placeholder 可以对上面定义的模型进行训练，但现在我们的目标是全部使用 tf.contrib.slim 来进行托管式的训练。要完成这个目标，需要借助两个函数的辅助，分别是上一篇文章 TensorFlow 自定义生成 .record 文件 定义的 get_record_dataset 函数和 slim 模块封装的 slim.learning.train 函数，前者用于获取训练数据，后者则执行对模型的训练。以下，先将训练用的代码列举出来（命名为 train.py）：
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Mar 30 19:27:44 2018
@author: shirhe-lyh
"""
"""Train a CNN model to classifying 10 digits.
Example Usage:
---------------
python3 train.py \
--record_path: Path to training tfrecord file.
--logdir: Path to log directory.
"""
import tensorflow as tf
import model
slim = tf.contrib.slim
flags = tf.app.flags
flags.DEFINE_string('record_path', None, 'Path to training tfrecord file.')
flags.DEFINE_string('logdir', None, 'Path to log directory.')
FLAGS = flags.FLAGS
def get_record_dataset(record_path,
reader=None, image_shape=[28, 28, 3],
num_samples=50000, num_classes=10):
"""Get a tensorflow record file.
Args:
"""
if not reader:
reader = tf.TFRecordReader
keys_to_features = {
'image/encoded':
tf.FixedLenFeature((), tf.string, default_value=''),
'image/format':
tf.FixedLenFeature((), tf.string, default_value='jpeg'),
'image/class/label':
tf.FixedLenFeature([1], tf.int64, default_value=tf.zeros([1],
dtype=tf.int64))}
items_to_handlers = {
'image': slim.tfexample_decoder.Image(shape=image_shape,
#image_key='image/encoded',
#format_key='image/format',
channels=3),
'label': slim.tfexample_decoder.Tensor('image/class/label', shape=[])}
decoder = slim.tfexample_decoder.TFExampleDecoder(
keys_to_features, items_to_handlers)
labels_to_names = None
items_to_descriptions = {
'image': 'An image with shape image_shape.',
'label': 'A single integer between 0 and 9.'}
return slim.dataset.Dataset(
data_sources=record_path,
reader=reader,
decoder=decoder,
num_samples=num_samples,
num_classes=num_classes,
items_to_descriptions=items_to_descriptions,
labels_to_names=labels_to_names)
def main(_):
dataset = get_record_dataset(FLAGS.record_path)
data_provider = slim.dataset_data_provider.DatasetDataProvider(dataset)
image, label = data_provider.get(['image', 'label'])
inputs, labels = tf.train.batch([image, label],
batch_size=64,
allow_smaller_final_batch=True)
cls_model = model.Model(is_training=True, num_classes=10)
preprocessed_inputs = cls_model.preprocess(inputs)
prediction_dict = cls_model.predict(preprocessed_inputs)
loss_dict = cls_model.loss(prediction_dict, labels)
loss = loss_dict['loss']
postprocessed_dict = cls_model.postprocess(prediction_dict)
classes = postprocessed_dict['classes']
acc = tf.reduce_mean(tf.cast(tf.equal(classes, labels), 'float'))
tf.summary.scalar('loss', loss)
tf.summary.scalar('accuracy', acc)
optimizer = tf.train.MomentumOptimizer(learning_rate=0.01, momentum=0.9)
train_op = slim.learning.create_train_op(loss, optimizer,
summarize_gradients=True)
slim.learning.train(train_op=train_op, logdir=FLAGS.logdir,
save_summaries_secs=20, save_interval_secs=120)
if __name__ == '__main__':
tf.app.run()
进行训练时，我们先根据上一篇文章的方式将所有训练图像（以及相应的标签）写成 tfrecord 文件，这一步如果训练数据足够多的话会耗费较长时间，但可以方便后续的数据读取（读取数据很快）。然后，我们接着上一篇文章的后一部分，要把这些数据读出来喂给模型。我们已经定义好了函数 get_record_dataset，它读取 .record 文件并返回一个 slim.dataset.Dataset 类对象。此时，用 slim.dataset_data_provider.DatasetDataProvider 作用一下，便可以用该类的函数 .get 来方便的将数据并行的提取出来（见上一篇文章）。因为并行的缘故，它每次返回的是单张图像，这时你可以对图像进行非批量的预处理，也可以直接使用 tf.train.batch 来生成指定大小的一个批量然后进行批量预处理。
一旦将训练数据提取出来，我们就可以把它们传给模型了（如上述代码 main 函数中间片段），最后的两句 tf.summary.sccalar 表示把损失和准确率写入到训练日志文件，方便之后我们在 tensorboard 中观察损失、准确率曲线。然后，再声明使用动量的随机梯度下降优化算法，紧接着便来到了训练的最后一步：
slim.learning.train(train_op=train_op, logdir=FLAGS.logdir,
save_summaries_secs=20, save_interval_secs=True)
这个函数将所有训练的迭代过程都封装起来，包括日志文件书写和模型保存。函数
slim.learning.train(train_op, logdir, train_step_fn=train_step,
train_step_kwargs=_USE_DEFAULT,
log_every_n_steps=1, graph=None, master='',
is_chief=True, global_step=None,
number_of_steps=None, init_op=_USE_DEFAULT,
init_feed_dict=None, local_init_op=_USE_DEFAULT,
init_fn=None, ready_op=_USE_DEFAULT,
summary_op=_USE_DEFAULT,
save_summaried_secs=600,
summary_writer=_USE_DEFAULT,
startup_delay_steps=0, saver=None,
save_interval_secs=600, sync_optimizer=None,
session_config=None, session_wrapper=None,
trace_every_n_steps=None,
ignore_live_threads=False)
参数众多，其中重要的有：1.train_op，指定优化算法；2.logdir，指定训练数据保存文件夹；3.save_summaries_secs，指定每隔多少秒更新一次日志文件（对应 tensorboard 刷新一次的时间）；4.save_interval_secs，指定每隔多少秒保存一次模型。
注意这段代码并没有定义数据传入的占位符，因此模型训练完成之后，我们并不知道怎么用，囧。不过，没关系，我们先训练完再说，在终端执行：
python3 train.py --record_path path/to/.record --logdir path_to_log_directory
会在 logdir 指定的路径下生成多个训练文件，而且每隔 save_interval_secs 会自动更新这些文件。当觉得模型训练到可以终止的时候，可以使用 Ctrl + C 来强制终止，或者通过指定参数 number_of_steps 来终止。而当觉得有必要对所训练的模型继续进行训练时，重新执行上述命令即可。如果你想查看训练的损失和准确率变化情况，在终端使用命令：
tensorboard --logdir path_to_log_directory
即可，至于查看的时间可以在训练过程中，也可以在训练结束后。
好了，有关使用 tf.contrib.slim 模块来进行模型训练的内容就讲完了，接下来还要解决一个重要的问题：没有定义占位符，我们怎么调用模型进行推断？
二、自定义模型导出（.pb 格式）
上面训练的模型会在 logdir 目录下保存为 .ckpt 格式的文件， 但一个致命的问题是没有数据入口，不知如何用它来对图像进行分类。为此，需要人为的添加数据入口、出口的结点。
要达到这个目的，我们需要仔细的研究一下下面这个用于模型导出的文件：
然后基于该文件做部分修改，用来导出我们的模型。请期待下一篇文章，我们将完成这一任务。
预告：下一篇文章将详细说明自定义的将 .ckpt 文件导出为 .pb 文件，敬请期待。
