在训练中使用 tf.dataset API 无法获取全部数据

Question

我在使用 GPU 进行训练时使用 tf.dataset 来获取图像、标签和边缘。但我发现数据集 API 无法加载所有数据。 我使用代码：

def get_dataset(filenames, shuffle_buffer, repeat_times, batch_size):
    dataset = tf.data.TFRecordDataset([filenames])
    dataset = dataset.map(tfrecord_preprocess)
    if repeat_times is None:
        dataset = dataset.repeat()
    else:
        dataset = dataset.repeat(repeat_times)
    dataset = dataset.shuffle(shuffle_buffer).batch(batch_size)
    return dataset

def tfrecord_preprocess(example):
    feature = {"image": tf.FixedLenFeature((), tf.string, default_value=""),
               "label": tf.FixedLenFeature((), tf.string, default_value=""),
               "edge": tf.FixedLenFeature((), tf.string, default_value="")}
    parsed_feature = tf.parse_single_example(example, feature)

    image = tf.decode_raw(parsed_feature["image"], out_type=tf.uint8)
    label = tf.decode_raw(parsed_feature["label"], out_type=tf.uint8)
    edge = tf.decode_raw(parsed_feature["edge"], out_type=tf.uint8)

    image = tf.cast(tf.reshape(image, shape=[1, 128, 128]), tf.float32) 
    label = tf.cast(tf.reshape(label, shape=[1, 128, 128]), tf.float32)
    edge = tf.cast(tf.reshape(edge, shape=[128, 128]), tf.float32)
    return image, label, edge

我写了一个简单的代码来测试API

dataset = get_dataset(filenames, shuffle_buffer, repeat_times, batchsize)
#shuffle=1000, repeat_times=2, batchsize=13
iter = dataset.make_one_shot_iterator
images, labels, edges = iter.get_next()
count = 0
with tf.Session() as sess:
    for _ in xrange(40):
        try:
            edges_value = sess.run(edges)
            count = count+len(edges_value)
            print count
        except tf.errors.OutofRangeError:
            break

数据的数量是260，所以在重复和批处理之后，epochs应该是40。它有效。

但是，当我使用类似代码进行训练时，数据总数少于 260，只有 140（通过 var 计数）。有谁知道如何解决这个问题？请帮助我。

我使用 tensorflow-gpu 1.4

我的训练代码是：

shuffle_buffer = params["shuffle_buffer"] #1000
repeat_times = params["repeat_times"] #1
batch_size = params["batch_size"] #26
num_classes = params["num_classes"] #2

dataset = model.get_dataset(filenames, shuffle_buffer, repeat_times, batch_size)
iterator = dataset.make_one_shot_iterator()
with tf.device('/gpu:1'): 
    global_step = tf.train.get_or_create_global_step()    
    learning_rate = tf.train.exponential_decay(params["learning_rate"], 
                                           global_step, 100, 0.99)    
    optimizer = tf.train.AdamOptimizer(learning_rate)

    images, labels, edges = iterator.get_next()
    _, probs = model.interence(features=images, training=True)
    loss, reg = model.get_loss(probs, labels, edges, num_classes)
    _, acc_mean, _ = model.get_acc(probs, labels)

    train_op = optimizer.minimize(loss, global_step=global_step)

    variables_average = tf.train.ExponentialMovingAverage(0.99, global_step)
    var_list = tf.trainable_variables(scope='.*(kernel|bias)')
    variables_average_op = variables_average.apply(var_list)    

    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

    with tf.control_dependencies(update_ops):
        train_all_op = tf.group(train_op, variables_average_op)

tf.summary.scalar("loss", loss)
tf.summary.scalar("reg", reg)
tf.summary.scalar("acc_mean", acc_mean)

merged = tf.summary.merge_all()
saver = tf.train.Saver(max_to_keep=5)

config = tf.ConfigProto(log_device_placement=True,
                        allow_soft_placement=True)
config.gpu_options.allow_growth = True
count = 0
with tf.Session(config=config) as sess:
    tf.global_variables_initializer().run() 
    writer = tf.summary.FileWriter('./train', sess.graph)
    for _ in xrange(10):
        try:
            edges_value = sess.run(edges)
            count = count+len(edges_value)
            _, step, summary = sess.run([train_all_op, global_step, merged])
            writer.add_summary(summary, step)
            if step % 5 == 0:
                loss_value = sess.run(loss)
                print loss_value
                acc_mean_value = sess.run(acc_mean)
                print acc_mean_value
                saver.save(sess, params["save_dir"], step)
        except tf.errors.OutOfRangeError:
            print "end of data"
            break
    print count
    print "the final step is %d" % step
    loss_value = sess.run(loss)
    print loss_value
    acc_mean_value = sess.run(acc_mean)
    print acc_mean_value
    saver.save(sess, params["save_dir"], step)    
    writer.close()

我终于在终端得到了信息：

end of data
130
the final step is 5

为了测试我将重复次数设置为 1 的代码

但我使用测试代码：

def test():
    dataset = get_dataset("train_output.tfrecords", 1000, 1, 26)
    terator = dataset.make_one_shot_iterator()
    images, labels, edges = iterator.get_next()

    count = 0
    with tf.Session() as sess:
        for i in xrange(10):
            try:
                images_value, labels_value, edges_value = sess.run([images, labels, edges])
                count = count+len(edges_value)
            except tf.errors.OutOfRangeError:
                print "end of data"

        print count
        print i

test()

终端显示：

260
9

Answer 1

问题是sess.run(edges) 导致这部分图表再次执行：images, labels, edges = iterator.get_next()。因此，每次运行它时，您都会消耗一次未计入计数器的迭代。

要获取边数，请在 with tf.device('/gpu:1') 块内保留一个计数器。您甚至可以使用 tf.summary.scalar 在 tensorboard 上绘制它，类似于使用 loss 的方式。

声明edges_count = tf.Variable(1, name='edges_count', trainable=False, dtype=tf.int32)

images, labels, edges = iterator.get_next()
edges_count_update_op = tf.assign_add(edges_count, len(edges))

然后将 edges_count_update_op 添加到您的 train_op 组。