TensorflowのSupervisorの使い方


tensorflowにはsupervisorという便利なツールがある．

- チェックポイントを自動で保存，そこから再開するから，1週間かかる学習も安心
- queueを自動でスタート，大量のtfrecordファイルも読み込める
- summaryも自動保存でTensorboardもバッチリ

ということで大規模学習に最適なのに，使い方のサンプルがほとんどみつからない．
そこでminimal sampleを作ってみた．


# スクリプト

モデルは$y=ax+b$というシンプルな線形回帰．

## データ生成

まずはtfrecordファイルを作るデータ生成スクリプト．

```python:generator.py
import tensorflow as tf
import numpy as np
import os.path

# parameters
a = 5
b = -3.2

# input
x_train = np.arange(0, 5, 0.1)
x_test  = np.arange(5, 10, 0.1)

# output
y_train = a * x_train + b + np.random.normal(scale=0.01, size=len(x_train))
y_test  = a * x_test  + b + np.random.normal(scale=0.01, size=len(x_test))



def save_data(x, y, s='train'):
    '''
    save input x and output y to a single tfrecord file
    '''
    def _bytes_feature(value):
        return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

    for i in range(len(x)):
        tfrecords_filename = os.path.join('./', '{:03d}_{}.tfrecords'.format(i, s))
        writer = tf.python_io.TFRecordWriter(tfrecords_filename)

        example = tf.train.Example(features=tf.train.Features(feature={
            'x': _bytes_feature(np.array([x[i]]).astype(np.float32).tostring()),
            'y': _bytes_feature(np.array([y[i]]).astype(np.float32).tostring()),
            })
        )
        writer.write(example.SerializeToString())
        writer.close()


save_data(x_train, y_train, 'train')
save_data(x_test, y_test, 'test')
```

## 学習・テスト

そして学習．オプションでテストにもなる．

```python:train.py

import argparse
import sys
import time
from datetime import datetime

import tensorflow as tf
from tensorflow import gfile
from tensorflow import logging



def read_and_decode(filename_queue):

    reader = tf.TFRecordReader()
    key, value = reader.read(filename_queue)

    features = tf.parse_single_example(
        value,
        features={
            'x': tf.FixedLenFeature([], tf.string),
            'y': tf.FixedLenFeature([], tf.string),
        })

    x = tf.decode_raw(features['x'], tf.float32)
    x = tf.reshape(x, [1])
    y = tf.decode_raw(features['y'], tf.float32)
    y = tf.reshape(y, [1])

    return x, y


def model(x):
    a = tf.Variable(tf.truncated_normal([1, 1], stddev=0.1))
    b = tf.Variable(tf.constant(0.1, shape=[1]))
    y = a * x + b
    return y, a, b


def load_batch():

    files = gfile.Glob(FLAGS.train_tfrecords)
    if not files:
        raise IOError("Can't find files: " + FLAGS.train_tfrecords)
    print(files)

    filename_queue = tf.train.string_input_producer(files,
                                                    num_epochs=FLAGS.num_epochs,
                                                    shuffle=FLAGS.train) # shuffle only for training

    x, y = read_and_decode(filename_queue)

    x_batch, y_batch = tf.train.batch([x, y], batch_size=FLAGS.batch_size)

    return x_batch, y_batch


def main(_):

    with tf.Graph().as_default():
    
        # data
        x_batch, y_batch = load_batch()

        # model
        y_conv, a, b = model(x_batch)

        # loss
        rmse = tf.sqrt(tf.losses.mean_squared_error(y_batch, y_conv))
        tf.summary.scalar('rmse', rmse)

        # minimizer
        global_step = tf.Variable(0, trainable=False)
        train_step = tf.train.AdamOptimizer(0.01).minimize(rmse, global_step=global_step)

        # supervisor
        if FLAGS.train:
            save_model_secs = 1 # save checkpoint every one second for training
        else:
            save_model_secs = 0 # do not save checkpoint

        sv = tf.train.Supervisor(logdir=FLAGS.checkpoint_dir,
                                 global_step=global_step,
                                 save_summaries_secs=1, # save summary every one second
                                 save_model_secs=save_model_secs,
                                 )

        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True # always useful

        with sv.managed_session(config=config) as sess:

            # training loop
            try:
                step = 0
                while not sv.should_stop():

                    if FLAGS.train:
                        sess.run([train_step])

                    x_input, y_est, y_true = sess.run([x_batch, y_batch, y_conv])
                    print('x={:.3f}, y_est={:.3f}, y_true={:.3f}'.format(x_input[0][0], y_est[0][0], y_true[0][0]))
    
                    if step % 100 == 0 and FLAGS.train:
                        rmse_value, g_step, a_val, b_val = sess.run([rmse, global_step, a, b])
                        print('global step {:05d} '.format(g_step), end=" ")
                        print('rmse {:010.2f} '.format(rmse_value), end=" ")
                        print(a_val[0][0], end=" ") # 1x1 matrix
                        print(b_val[0])

                    step += 1
    
            except tf.errors.OutOfRangeError:
                print('Training ends with {} steps'.format(step))
    
            sv.Stop()
            
        print('end')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--train', dest='train', action='store_true')
    parser.add_argument('--test', dest='train', action='store_false')
    parser.set_defaults(train=True)

    parser.add_argument(
        '--train_tfrecords',
        type=str,
        default='./*.tfrecords',
        help='regrep pattern of tfrecord files'
    )
    parser.add_argument(
        '--num_epochs',
        type=int,
        default=100,
        help='Number of epochs'
    )
    parser.add_argument(
        '--batch_size',
        type=int,
        default=1,
        help='batch size'
    )
    parser.add_argument(
        '--checkpoint_dir',
        type=str,
        default='./tmp',
        help='directory where checkpoint files are stored'
    )
    FLAGS, unparsed = parser.parse_known_args()
    tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)

```

# 使い方

## まずはデータを生成．

*.tfrecordsというファイルがたくさんできる．

```bash:データ生成
$ python3 generator.py
```

## では学習．

```bash:学習
$ python3 train.py --train_tfrecords=*_train.tfrecords

['./000_train.tfrecords', './001_train.tfrecords', './002_train.tfrecords', './003_train.tfrecords', './004_train.tfrecords', './005_train.tfrecords', './006_train.tfrecords',
...
'./041_train.tfrecords', './042_train.tfrecords', './043_train.tfrecords', './044_train.tfrecords', './045_train.tfrecords', './046_train.tfrecords', './047_train.tfrecords', './048_train.tfrecords', './049_train.tfrecords']
x=3.500, y_est=14.303, y_true=-0.477
global step 00001  rmse 0000013.74  -0.162092 0.09
x=3.700, y_est=15.290, y_true=-0.497
x=1.200, y_est=2.803, y_true=-0.088
x=1.300, y_est=3.311, y_true=-0.087
x=0.200, y_est=-2.200, y_true=0.080
x=3.900, y_est=16.289, y_true=-0.379
x=3.800, y_est=15.809, y_true=-0.334
x=3.000, y_est=11.810, y_true=-0.209
x=2.500, y_est=9.307, y_true=-0.126
x=0.700, y_est=0.296, y_true=0.070
x=4.900, y_est=21.311, y_true=-0.278
x=2.700, y_est=10.295, y_true=-0.057
x=4.300, y_est=18.293, y_true=-0.133
x=1.000, y_est=1.800, y_true=0.108
x=4.000, y_est=16.790, y_true=-0.025
x=2.200, y_est=7.809, y_true=0.093
x=1.700, y_est=5.301, y_true=0.137
x=0.900, y_est=1.300, y_true=0.180
x=2.300, y_est=8.290, y_true=0.173
...
global step 02401  rmse 0000000.02  4.99713 -3.18101
...
x=1.100, y_est=2.313, y_true=2.303
x=3.900, y_est=16.289, y_true=16.301
x=4.900, y_est=21.311, y_true=21.301
x=4.500, y_est=19.292, y_true=19.303
x=0.800, y_est=0.804, y_true=0.804
x=1.800, y_est=5.797, y_true=5.798
x=2.100, y_est=7.292, y_true=7.290
x=3.500, y_est=14.303, y_true=14.272
x=3.700, y_est=15.290, y_true=15.267
x=1.000, y_est=1.800, y_true=1.793
x=4.600, y_est=19.804, y_true=19.760
x=1.600, y_est=4.792, y_true=4.796
x=1.900, y_est=6.305, y_true=6.303
x=1.400, y_est=3.805, y_true=3.815
x=4.000, y_est=16.790, y_true=16.819
x=4.200, y_est=17.814, y_true=17.826
x=3.800, y_est=15.809, y_true=15.831
x=2.700, y_est=10.295, y_true=10.328
x=0.700, y_est=0.296, y_true=0.321
x=2.200, y_est=7.809, y_true=7.815
x=2.000, y_est=6.790, y_true=6.807
Training ends with 2475 steps
end
```

学習データはシャッフルされてバッチになる．

supervisorがチェックポイントを保存しているので，もう一回学習すると保存済みチェックポイントから学習を再開する．



```bash:もう一回学習
$ python3 train.py --train_tfrecords=*_train.tfrecords

['./000_train.tfrecords', './001_train.tfrecords', './002_train.tfrecords', './003_train.tfrecords', './004_train.tfrecords', './005_train.tfrecords', './006_train.tfrecords', 
...
'./044_train.tfrecords', './045_train.tfrecords', './046_train.tfrecords', './047_train.tfrecords', './048_train.tfrecords', './049_train.tfrecords']
x=2.000, y_est=6.790, y_true=6.820
global step 02208  rmse 0000000.02  5.00433 -3.18855
x=0.300, y_est=-1.696, y_true=-1.690
x=3.900, y_est=16.289, y_true=16.302
x=0.600, y_est=-0.178, y_true=-0.201
x=3.300, y_est=13.299, y_true=13.274
x=0.200, y_est=-2.200, y_true=-2.203
x=3.600, y_est=14.800, y_true=14.767
x=0.400, y_est=-1.193, y_true=-1.204
x=0.800, y_est=0.804, y_true=0.798
...
global step 04608  rmse 0000000.02  5.00579 -3.20404
...
x=3.300, y_est=13.299, y_true=13.295
x=3.000, y_est=11.810, y_true=11.809
x=0.500, y_est=-0.684, y_true=-0.688
x=2.600, y_est=9.818, y_true=9.820
x=0.700, y_est=0.296, y_true=0.315
x=3.400, y_est=13.796, y_true=13.822
x=2.900, y_est=11.314, y_true=11.317
x=1.100, y_est=2.313, y_true=2.309
x=3.600, y_est=14.800, y_true=14.801
x=2.800, y_est=10.811, y_true=10.800
x=2.400, y_est=8.811, y_true=8.796
x=1.500, y_est=4.308, y_true=4.298
x=3.800, y_est=15.809, y_true=15.799
Training ends with 2471 steps
end
```



## そしてテスト

テストのバッチはシャッフルしないので（string_input_producerのshuffleがFalse），入力xが先頭から順番に評価される．


```bash:そしてテスト
$ python3 train.py --train_tfrecords=*_test.tfrecords --test --num_epochs=1

['./000_test.tfrecords', './001_test.tfrecords', './002_test.tfrecords', './003_test.tfrecords', './004_test.tfrecords', './005_test.tfrecords', './006_test.tfrecords', './007_test.tfrecords', 
...
'./042_test.tfrecords', './043_test.tfrecords', './044_test.tfrecords', './045_test.tfrecords', './046_test.tfrecords', './047_test.tfrecords', './048_test.tfrecords', './049_test.tfrecords']
x=5.100, y_est=22.314, y_true=22.277
x=5.200, y_est=22.800, y_true=22.776
x=5.300, y_est=23.312, y_true=23.276
x=5.400, y_est=23.798, y_true=23.775
x=5.500, y_est=24.306, y_true=24.274
x=5.600, y_est=24.805, y_true=24.774
x=5.700, y_est=25.309, y_true=25.273
x=5.800, y_est=25.792, y_true=25.773
x=5.900, y_est=26.306, y_true=26.272
x=6.000, y_est=26.796, y_true=26.772
x=6.100, y_est=27.303, y_true=27.271
x=6.200, y_est=27.802, y_true=27.771
x=6.300, y_est=28.311, y_true=28.270
x=6.400, y_est=28.785, y_true=28.770
x=6.500, y_est=29.285, y_true=29.269
x=6.600, y_est=29.786, y_true=29.769
x=6.700, y_est=30.304, y_true=30.268
x=6.800, y_est=30.794, y_true=30.768
x=6.900, y_est=31.296, y_true=31.267
x=7.000, y_est=31.792, y_true=31.767
x=7.100, y_est=32.321, y_true=32.266
x=7.200, y_est=32.806, y_true=32.766
x=7.300, y_est=33.288, y_true=33.265
x=7.400, y_est=33.802, y_true=33.765
x=7.500, y_est=34.301, y_true=34.264
x=7.600, y_est=34.788, y_true=34.764
x=7.700, y_est=35.300, y_true=35.263
x=7.800, y_est=35.793, y_true=35.763
x=7.900, y_est=36.291, y_true=36.262
x=8.000, y_est=36.822, y_true=36.761
x=8.100, y_est=37.295, y_true=37.261
x=8.200, y_est=37.801, y_true=37.760
x=8.300, y_est=38.304, y_true=38.260
x=8.400, y_est=38.799, y_true=38.759
x=8.500, y_est=39.311, y_true=39.259
x=8.600, y_est=39.797, y_true=39.758
x=8.700, y_est=40.318, y_true=40.258
x=8.800, y_est=40.795, y_true=40.757
x=8.900, y_est=41.314, y_true=41.257
x=9.000, y_est=41.790, y_true=41.756
x=9.100, y_est=42.297, y_true=42.256
x=9.200, y_est=42.789, y_true=42.755
x=9.300, y_est=43.297, y_true=43.255
x=9.400, y_est=43.803, y_true=43.754
x=9.500, y_est=44.311, y_true=44.254
x=9.600, y_est=44.805, y_true=44.753
x=9.700, y_est=45.296, y_true=45.253
x=9.800, y_est=45.796, y_true=45.752
x=9.900, y_est=46.303, y_true=46.252
Training ends with 49 steps
end
```