Skip to content

secsilm/understaing-datasets-estimators-tfrecords

Repository files navigation

Estimators 和 Datasets

Google 在 2017 年 9 月 12 号的博文 Introduction to TensorFlow Datasets and Estimators 中介绍了新引入的两个新特性 DatasetsEstimators

  • Datasets:创建一个输入管道(input pipelines)来为你的模型读取数据,在这个 pipelines 中你可以做一些数据预处理,尽量都使用 TensorFlow 自己的函数,即 tf 开头的函数(比如 tf.reshape),这样可以提高程序执行效率。

  • Estimators:这是模型的核心部分,而 Estimators 的核心部分则是一个 model_fn 函数(后面会细讲),你在这个函数中定义你的模型架构,输入是特征和标签,输出是一个定义好的 estimator。

tensorflow architecture
TensorFlow 架构,图自 Google Developers Blog

实际上这两个特性并不是第一次引入,只不过之前是放在 tf.contrib 里,而这次是引入到了 TensorFlow 核心组件中,意味着可以在生产环境中使用。我 6 月份的时候也写过一篇博文简单说了下 tf.contrib.learn.DNNRegressor 的使用,实际上这就是 Estimators 内置的一个模型(estimator)。这两个都是高层 API,也就是说为了创建一个模型你不用再写一些很底层的代码(比如定义权重偏置项),可以像 scikit-learn 和 Keras 那样很轻松的几行代码创建一个模型,便于快速实现。

本篇博文就是试图将这两个高层 API 结合起来,使用 TensorFlow 的数据格式 TFRecords 来实现一个在 CIFAR-10 数据集上的 CNN 模型。完整代码可在我的 GitHub 上找到。

Note:本篇博文中的模型并不是结果最好的模型,仅仅是为了展示如何将 Estimators 和 Datasets 结合起来使用。

更新

我会在这里列出对本文的更新。

  • 2018 年 5 月 1 日:增加使用已训练的模型进行预测的 demo。

用法

你可以使用 python cifar10_estimator_dataset.py --help 来查看可选参数:

       USAGE: cifar10_estimator_dataset.py [flags]
flags:

cifar10_estimator_dataset.py:
  --batch_size: Batch size
    (default: '64')
    (an integer)
  --dropout_rate: Dropout rate
    (default: '0.5')
    (a number)
  --eval_dataset: Filename of evaluation dataset
    (default: 'eval.tfrecords')
  --learning_rate: Learning rate
    (default: '0.001')
    (a number)
  --model_dir: Filename of testing dataset
    (default: 'models/cifar10_cnn_model')
  --num_epochs: Number of training epochs
    (default: '10')
    (an integer)
  --test_dataset: Filename of testing dataset
    (default: 'test.tfrecords')
  --train_dataset: Filename of training dataset
    (default: 'train.tfrecords')

TFRecords 和 TensorBoard 文件(包括我做的所有 run)较大,没有放到 GitHub 上,你可以从百度盘上获取:

模型架构

为了让大家对模型架构先有个清晰地了解,我先把 TensorBoard (不熟悉 TensorBoard 的话可以参考这里)中显示的模型架构图贴出来(数据集我也就不介绍了,这是个很常用的数据集,如有不熟悉的可以参看这里):

cnn model
模型架构

可以看到两层卷积层,两层池化层,两层 BN 层,一层 dropout,三层全连接层(DENSE)。

读取数据集

在给模型「喂」数据的时候,我们的流程大概是这样的:

  1. 创建一个 Dataset 对象来表示我们的数据集,有多种方法可以创建一个 Dataset 对象,我说几个比较常用的:

    • tf.data.Dataset.from_tensor_slices():这种方法适合于你的数据集是 numpy 数组类型的。
    • tf.data.TFRecordDataset()这是本文所使用的方法,适合于你的数据集是 TFRecords 类型的。
    • tf.data.TextLineDataset():适合于你的数据集是 txt 格式的。
  2. 对数据集进行一些预处理:

    • Dataset.map():和普通的 map 函数一样,对数据集进行一些变换,例如图像数据集的类型转换(uint8 -> float32)以及 reshape 等。
    • Dataset.shuffle():打乱数据集
    • Dataset.batch():将数据集切分为特定大小的 batch
    • Dataset.repeat():将数据集重复多次。如果不使用这个方法,在第一次遍历到数据集的结尾的时候,会抛出一个 tf.errors.OutOfRangeError 异常,表示数据集已经遍历完毕。但是实际中我们可能需要对数据集迭代训练不止一次,这时候就要用 repeat() 来重复数据集多次。如果不加任何参数,那么表示重复数据集无穷多次。
  3. 使用 Iteratorget_next() 方法来每次获取一个 batch 的数据(假如你是使用 mini-batch 训练的话)。目前 TensorFlow 提供四种 Iterator(详细见 Creating an iterator):

    • one-shot这是本文程序所使用的方法,使用 Dataset.make_one_shot_iterator() 来创建,不需要初始化。官方有这么一句话:Note: Currently, one-shot iterators are the only type that is easily usable with an Estimator. 不过呢,我也发现外国友人 Peter Roelants 写了个例子将下面的 initializable Iterator 和 Estimator 一起使用,见 Example using TensorFlow Estimator, Experiment & Dataset on MNIST data
    • initializable:使用 Dataset.make_initializable_iterator() 创建,需要使用 iterator.initializer 初始化。该方式可以允许你自定义数据集,例如你的数据集是 range(0, max_value),这里面 max_value 是一个 Tensor,在初始化的时候你需要赋值。
    • reinitializable:这是种比较复杂的方式,简单来说也就是使你可以从多个不同的 Dataset 对象获取数据,详细可见 Creating an iterator
    • feedable:同样比较复杂,当然更灵活,可以针对不同的 Dataset 对象和 tf.Session.run 使用不同的 Iterator,详细可见 Creating an iterator

在 Estimator 中,我们输入必须是一个函数,这个函数必须返回特征和标签(或者只有特征),所以我们需要把上面的内容写到一个函数中。因为训练输入和验证输入是不一样的,所以需要两个输入函数:train_input_fneval_input_fn。为了保持文章简洁,我下面只列出 train_input_fneval_input_fn 和其大同小异。

此处我使用了 tf.data.TFRecordDataset,所以你需要将你的数据集写成 TFRecords 格式,比如 train.tfrecords。TFRecords 格式每行表示一个样本(record),关于如何将数据集写成 TFRecords 格式,我将在另一篇博文中说明。

def train_input_fn():
    '''
    训练输入函数,返回一个 batch 的 features 和 labels
    '''
    train_dataset = tf.data.TFRecordDataset(FLAGS.train_dataset)
    train_dataset = train_dataset.map(parser)
    # num_epochs 为整个数据集的迭代次数
    train_dataset = train_dataset.repeat(FLAGS.num_epochs)
    train_dataset = train_dataset.batch(FLAGS.batch_size)
    train_iterator = train_dataset.make_one_shot_iterator()

    features, labels = train_iterator.get_next()
    return features, labels

而其中的 map 函数的参数 parser 也是一个函数,用于将图片和标签从 TFRecords 中解析出来。

def parser(record):
    keys_to_features = {
        'image_raw': tf.FixedLenFeature((), tf.string),
        'label': tf.FixedLenFeature((), tf.int64)
    }
    parsed = tf.parse_single_example(record, keys_to_features)
    image = tf.decode_raw(parsed['image_raw'], tf.uint8)
    image = tf.cast(image, tf.float32)
    label = tf.cast(parsed['label'], tf.int32)
    return image, label

到此,关于模型的 input pipeline 就差不多结束了。下面就是模型的核心部分了:定义一个模型函数 model_fn

定义模型函数

上面是定义了 input pipeline,那么现在该来定义模型架构了。模型大致架构就是上面的模型架构图。该函数需要返回一个定义好的 tf.estimator.EstimatorSpec 对象,对于不同的 mode,所必须提供的参数是不一样的:

  • 训练模式,即 mode == tf.estimator.ModeKeys.TRAIN,必须提供的是 losstrain_op
  • 验证模式,即 mode == tf.estimator.ModeKeys.EVAL,必须提供的是 loss
  • 预测模式,即 mode == tf.estimator.ModeKeys.PREDICT,必须提供的是 predicitions

为保持文章简洁,我省略了一些重复性代码。

def cifar_model_fn(features, labels, mode):
    """Model function for cifar10 model"""
    # 输入层
    x = tf.reshape(features, [-1, 32, 32, 3])
    # 第一层卷积层
    x = tf.layers.conv2d(inputs=x, filters=64, kernel_size=[
                             3, 3], padding='same', activation=tf.nn.relu, name='CONV1')
    x = tf.layers.batch_normalization(
        inputs=x, training=mode == tf.estimator.ModeKeys.TRAIN, name='BN1')
    # 第一层池化层
    x = tf.layers.max_pooling2d(inputs=x, pool_size=[
                                    3, 3], strides=2, padding='same', name='POOL1')
    
    # 你可以添加更多的卷积层和池化层 ……

    # 全连接层
    x = tf.reshape(x, [-1, 8 * 8 * 128])
    x = tf.layers.dense(inputs=x, units=512, activation=tf.nn.relu, name='DENSE1')

    # 你可以添加更多的全连接层 ……

    logits = tf.layers.dense(inputs=x, units=10, name='FINAL')

    # 预测
    predictions = {
        'classes': tf.argmax(input=logits, axis=1, name='classes'),
        'probabilities': tf.nn.softmax(logits, name='softmax_tensor')
    }

    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

    # 计算损失(对于 TRAIN 和 EVAL 模式)
    onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth=10)
    loss = tf.losses.softmax_cross_entropy(onehot_labels, logits, scope='LOSS')
    
    # 评估方法
    accuracy, update_op = tf.metrics.accuracy(
        labels=labels, predictions=predictions['classes'], name='accuracy')
    batch_acc = tf.reduce_mean(tf.cast(
        tf.equal(tf.cast(labels, tf.int64), predictions['classes']), tf.float32))
    tf.summary.scalar('batch_acc', batch_acc)
    tf.summary.scalar('streaming_acc', update_op)

    # 训练配置(对于 TRAIN 模式)
    if mode == tf.estimator.ModeKeys.TRAIN:
        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        optimizer = tf.train.RMSPropOptimizer(learning_rate=FLAGS.learning_rate)
        with tf.control_dependencies(update_ops):
            train_op = optimizer.minimize(
                loss=loss, global_step=tf.train.get_global_step())
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

    eval_metric_ops = {
        'accuracy': (accuracy, update_op)
    }
    return tf.estimator.EstimatorSpec(mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)

至此,input pipeline 和模型都已经定义好了,下一步就是实际的 run 了。

Run

首先我们需要创建一个 tf.estimator.Estimator 对象:

cifar10_classifier = tf.estimator.Estimator(
        model_fn=cifar_model_fn, model_dir=FLAGS.model_dir)

其中 model_dir 是用于存放模型文件和 TensorBoard 文件的目录。

然后开始训练和验证:

cifar10_classifier.train(input_fn=train_input_fn)
eval_results = cifar10_classifier.evaluate(input_fn=eval_input_fn)

程序结束后你便可以在你的 model_dir 里看到类似如下的文件结构:

model_dir model_dir 中的文件结构

然后你可以使用 tensorboard --logdir=/your/model/dir(Linux 中你可能需要使用 python -m tensorboard.main --logdir=/your/model/dir)来在 TensorBoard 中查看训练信息,默认只有 SCALARSGRAPHS 面板是有效的,你也可以自己使用 tf.summary 来手动添加 summary 信息。

scalars SCALARS 面板

graphs GRAPHS 面板

使用训练好的模型进行预测

在训练好模型之后,模型文件已经保存到了 FLAGS.model_dir 中,那么在对新样本进行预测时只需要调用 estimator 的 predict() 方法进行预测就行了。

def infer(argv=None):
    '''Run the inference and return the result.'''
    config = tf.estimator.RunConfig()
    config = config.replace(model_dir=FLAGS.saved_model_dir)
    estimator = get_estimator(config)
    predict_input_fn = tf.estimator.inputs.numpy_input_fn(
        x=load_image(), shuffle=False)
    result = estimator.predict(input_fn=predict_input_fn)
    for r in result:
        print(r)


def load_image():
    '''Load image into numpy array.'''
    images = np.zeros((10, 3072), dtype='float32')
    for i, file in enumerate(Path('predict-images/').glob('*.png')):
        image = np.array(Image.open(file)).reshape(3072)
        images[i, :] = image
    return images

这是主要的两个函数,完整代码见 cifar10_estimator_dataset_predict.py

有几点需要说明:

  • 我把要预测的图片放在了 predict-images/ 文件夹下,你可以自由更改这个地址。
  • 这里我使用了 tf.estimator.inputs.numpy_input_fn() 来作为预测的输入函数,该函数可以直接接受 numpy array 作为输入。除此之外,你还可以像 cifar10_estimator_dataset.py 中的 train_input_fn() 一样,使用 tf.data.Dataset.from_tensor_slices() 或者 tf.data.TFRecordDataset(),再结合 Dataset.make_one_shot_iterator() 来定义一个预测输入函数。

用法很简单,假设你的模型文件放在 models/cifar10 下,那么在命令行执行下面的语句即可:

python cifar10_estimator_dataset_predict.py --saved_model_dir models/cifar10

--saved_model_dir 的默认值是 models/adam

执行完后可以看到类似下面这样的输出结果,当然下面的结果很差,由于时间有限我也没有过多的调模型,这里只是说明下过程:

predict

Summary

总的来说,使用 Datasets 和 Estimators 来训练模型大致就是这么几个步骤:

  1. 定义输入函数,在函数中对你的数据集做一些必要的预处理,返回 features 和 labels。
  2. 定义模型函数,返回 tf.estimator.EstimatorSpec 对象。
  3. 使用模型函数创建 tf.estimator.Estimator 对象。
  4. 使用创建好的对象 train and evaluate。

Notes

关于 num_epochs

如果你设置 num_epochs 为比如说 30,然而你在训练的时候看到类似如下的控制台输出:

INFO:tensorflow:global_step/sec: 0.476364
INFO:tensorflow:loss = 0.137512, step = 14901 (209.924 sec)
INFO:tensorflow:global_step/sec: 0.477139
INFO:tensorflow:loss = 0.0203241, step = 15001 (209.583 sec)
INFO:tensorflow:global_step/sec: 0.477511
INFO:tensorflow:loss = 0.132834, step = 15101 (209.419 sec)

你可以看到 step 已经上万了,这是因为这里的 step 指的是一个 batch 的训练迭代,而 num_epochs 设为 30 意味着你要把整个训练集遍历 30 次(也是我们通常的做法)。也就是说,假如你有 50000 个样本,batch 大小为 50,那么你的数据集将被切分为 1000 个 batch,也就是遍历一遍数据集需要 1000 step,所以说 num_epochs 为 30 时,你的程序需要到 step=30000 才会训练结束。所以切记 num_epochs 表示的是整个训练集的迭代次数。

References

  1. Introduction to TensorFlow Datasets and Estimators
  2. Importing Data | TensorFlow
  3. Creating Estimators in tf.estimator | TensorFlow
  4. Example using TensorFlow Estimator, Experiment & Dataset on MNIST data
  5. Higher-Level APIs in TensorFlow – Onfido Tech – Medium

END