【TensorFlow】理解 Estimators 和 Datasets

- **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](https://gist.github.com/peterroelants/9956ec93a07ca4e9ba5bc415b014bcca)。
- initializable：使用 `Dataset.make_initializable_iterator()` 创建，需要使用 `iterator.initializer` 初始化。该方式可以允许你自定义数据集，例如你的数据集是 `range(0, max_value)`，这里面 `max_value` 是一个 `Tensor`，在初始化的时候你需要赋值。
- reinitializable：这是种比较复杂的方式，简单来说也就是使你可以从多个不同的 `Dataset` 对象获取数据，详细可见 [Creating an iterator](https://www.tensorflow.org/programmers_guide/datasets#creating_an_iterator)。
- feedable：同样比较复杂，当然更灵活，可以针对不同的 `Dataset` 对象和 `tf.Session.run` 使用不同的 `Iterator`，详细可见 [Creating an iterator](https://www.tensorflow.org/programmers_guide/datasets#creating_an_iterator)。

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

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

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)

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

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

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)