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

目前，要在批处理张量上获得 fully-defined 静态形状，如果批处理大小不除总数，则需要明确告诉 TensorFlow "drop" 任何 "remainder"的元素均匀。为此，请替换以下行：

dataset = dataset.batch(opt.batchsize)

...应用 tf.contrib.data.batch_and_drop_remainder():

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

Tensorflow dataset.batch() 未显示实际批量大小

Tensorflow dataset.batch() not showing the real batch size

python

tensorflow

tensorflow-datasets