TensorFlow-slim 训练 CNN 分类模型(续)

        在前面的文章 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.scalar 表示把损失和准确率写入到训练日志文件,方便之后我们在 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 格式的文件, 但一个致命的问题是没有数据入口,不知如何用它来对图像进行分类。为此,需要人为的添加数据入口、出口的结点。

        要达到这个目的,我们需要仔细的研究一下下面这个用于模型导出的文件:

TensorFlow models/research/object_detection/export.py

然后基于该文件做部分修改,用来导出我们的模型。请期待下一篇文章,我们将完成这一任务。

预告:下一篇文章将详细说明自定义的将 .ckpt 文件导出为 .pb 文件,敬请期待。

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 214,837评论 6 496
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 91,551评论 3 389
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 160,417评论 0 350
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 57,448评论 1 288
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 66,524评论 6 386
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 50,554评论 1 293
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 39,569评论 3 414
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 38,316评论 0 270
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 44,766评论 1 307
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 37,077评论 2 330
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 39,240评论 1 343
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 34,912评论 5 338
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 40,560评论 3 322
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 31,176评论 0 21
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,425评论 1 268
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 47,114评论 2 366
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 44,114评论 2 352

推荐阅读更多精彩内容