tensorflow dataset.batch()未显示真实批次大小

Question

我想将基于队列的数据加载机制更改为 tf.data api。

原始代码是：

    # Index queue
    self.input_idxs = tf.placeholder(tf.int64, shape=[None, 2])
    idx_queue = tf.FIFOQueue(1e8, tf.int64)
    self.enq_idxs = idx_queue.enqueue_many(self.input_idxs)
    get_idx = idx_queue.dequeue()
    # Image loading queue
    img_queue = tf.FIFOQueue(opt.max_queue_size, task.proc_arg_dtype)
    load_data = tf.py_func(task.load_sample_data, [get_idx], task.proc_arg_dtype)
    enq_img = img_queue.enqueue(load_data)
    init_sample = img_queue.dequeue()
    # Preprocessing queue
    # (for any preprocessing that can be done with TF operations)
    data_queue = tf.FIFOQueue(opt.max_queue_size, task.data_arg_dtype,
                              shapes=task.data_shape)
    enq_data = data_queue.enqueue(task.preprocess(init_sample, train_flag))
    self.get_sample = data_queue.dequeue_many(opt.batchsize)

更改后，是：

    # Dataset
    self.input_idxs = tf.placeholder(tf.int64, shape=[None, 2])
    dataset = tf.data.Dataset.from_tensor_slices(self.input_idxs)
    def load_sample(idx):
        sample = task.load_sample_data(idx)
        sample = task.preprocess(sample, train_flag)
        return sample
    dataset = dataset.map(lambda idx: tf.py_func(load_sample, [idx], task.proc_arg_dtype), num_parallel_calls=self.num_threads)
    def gen(dataset):
        yield dataset.make_one_shot_iterator().get_next()
    dataset = tf.data.Dataset.from_generator(gen, tuple(task.proc_arg_dtype), tuple(task.data_shape))
    dataset = dataset.batch(opt.batchsize)
    self.iterator = dataset.make_initializable_iterator()
    self.get_sample = self.iterator.get_next()

其中 task.proc_arg_dtype和 task.data_shape是：

    proc_arg_dtype = [tf.float32, tf.float32, tf.int32, tf.int32, tf.int32, tf.float32, tf.int32, tf.int32, tf.int32]
    data_shape = [
        [opt.input_res, opt.input_res, 3],
        [opt.output_res, opt.output_res, opt.det_inputs],
        [2, opt.max_nodes, 2],
        [4],
        [opt.max_nodes, opt.obj_slots + opt.rel_slots],
        [opt.max_nodes, opt.obj_slots, 5],
        [opt.max_nodes, opt.rel_slots, 2],
        [opt.max_nodes, 7],
        [1]
    ]

由于我发现tf.py_func没有data_shape参数，因此我使用tf.data.Dataset.from_generator来执行此操作。(不确定这是正确的，因为我在竞争之前遇到了一个问题(

问题以前是self.get_sample类似于：

[<tf.Tensor 'IteratorGetNext:0' shape=(8, 512, 512, 3) dtype=float32>, <tf.Tensor 'IteratorGetNext:1' shape=(8, 64, 64, 300) dtype=float32>, <tf.Tensor 'IteratorGetNext:2' shape=(8, 2, 200, 2) dtype=int32>, <tf.Tensor 'IteratorGetNext:3' shape=(8, 4) dtype=int32>, <tf.Tensor 'IteratorGetNext:4' shape=(8, 200, 9) dtype=int32>, <tf.Tensor 'IteratorGetNext:5' shape=(8, 200, 3, 5) dtype=float32>, <tf.Tensor 'IteratorGetNext:6' shape=(8, 200, 6, 2) dtype=int32>, <tf.Tensor 'IteratorGetNext:7' shape=(8, 200, 7) dtype=int32>, <tf.Tensor 'IteratorGetNext:8' shape=(8, 1) dtype=int32>]

批处理大小是第一个维度。但是，通过使用dataset.batch(opt.batch_size)，self.get_sample是

[<tf.Tensor 'IteratorGetNext:0' shape=(?, 512, 512, 3) dtype=float32>, <tf.Tensor 'IteratorGetNext:1' shape=(?, 64, 64, 300) dtype=float32>, <tf.Tensor 'IteratorGetNext:2' shape=(?, 2, 200, 2) dtype=int32>, <tf.Tensor 'IteratorGetNext:3' shape=(?, 4) dtype=int32>, <tf.Tensor 'IteratorGetNext:4' shape=(?, 200, 9) dtype=int32>, <tf.Tensor 'IteratorGetNext:5' shape=(?, 200, 3, 5) dtype=float32>, <tf.Tensor 'IteratorGetNext:6' shape=(?, 200, 6, 2) dtype=int32>, <tf.Tensor 'IteratorGetNext:7' shape=(?, 200, 7) dtype=int32>, <tf.Tensor 'IteratorGetNext:8' shape=(?, 1) dtype=int32>]

未显示真实批处理大小。

Answer 1

当前，要在批处理张量上获得完全定义的静态形状，您需要明确地告诉TensorFlow如果批处理大小不划分总数均匀的元素。为此，请更换以下行：

dataset = dataset.batch(opt.batchsize)

...使用tf.contrib.data.batch_and_drop_remainder()的应用：

dataset = dataset.apply(tf.contrib.data.batch_and_drop_remainder(opt.batchsize))