如何在Tensorflow 2.0数据集中动态更改批量大小？

Question

在TensorFlow 1.X中，您可以使用占位符动态更改批量大小。例如： dataset.batch(batch_size=tf.placeholder()) 查看完整示例 在TensorFlow 2.0中如何实现呢？

我尝试了以下方法，但都不起作用。

import numpy as np
import tensorflow as tf


def new_gen_function():
    for i in range(100):
        yield np.ones(2).astype(np.float32)


batch_size = tf.Variable(5, trainable=False, dtype=tf.int64)
train_ds = tf.data.Dataset.from_generator(new_gen_function, output_types=(tf.float32)).batch(
    batch_size=batch_size)

for data in train_ds:
    print(data.shape[0])
    batch_size.assign(10)
    print(batch_size)

输出

5
<tf.Variable 'Variable:0' shape=() dtype=int64, numpy=10>
5
<tf.Variable 'Variable:0' shape=() dtype=int64, numpy=10>
5
<tf.Variable 'Variable:0' shape=() dtype=int64, numpy=10>
5
...
...

我正在使用Gradient tape自定义训练循环来训练模型。如何实现这一点？

- Himaprasoon

---

也许你可以尝试用tf.keras.Input替换tf.placeholder。你可以参考这个链接https://dev59.com/xVMH5IYBdhLWcg3wmwPw获取更多信息。谢谢！ - user11530462

Answer 1

我认为您不能像在TF1中那样做。

一个解决办法是通过堆叠单独的样本来构建批次：

import tensorflow as tf

ds = tf.data.Dataset.range(10).repeat()
iterator = iter(ds)
for batch_size in range(1, 10):
  batch = tf.stack([iterator.next() for _ in range(batch_size)], axis=0)
  print(batch)

# tf.Tensor([0], shape=(1,), dtype=int64)
# tf.Tensor([1 2], shape=(2,), dtype=int64)
# tf.Tensor([3 4 5], shape=(3,), dtype=int64)
# tf.Tensor([6 7 8 9], shape=(4,), dtype=int64)
# tf.Tensor([0 1 2 3 4], shape=(5,), dtype=int64)
# tf.Tensor([5 6 7 8 9 0], shape=(6,), dtype=int64)
# tf.Tensor([1 2 3 4 5 6 7], shape=(7,), dtype=int64)
# tf.Tensor([8 9 0 1 2 3 4 5], shape=(8,), dtype=int64)
# tf.Tensor([6 7 8 9 0 1 2 3 4], shape=(9,), dtype=int64)

- P-Gn

Answer 2

据我所知，您应该实例化一个新的数据集迭代器以使您的更改生效。这将需要进行一些微调以跳过已经看到的样本。

以下是我的最简单的解决方案：

import numpy as np
import tensorflow as tf

def get_dataset(batch_size, num_samples_seen):
    return tf.data.Dataset.range(
        100
    ).skip(
        num_samples_seen
    ).batch(
        batch_size=batch_size
    )

def main():
    batch_size = 1
    num_samples_seen = 0

    train_ds = get_dataset(batch_size, num_samples_seen)

    ds_iterator = iter(train_ds)
    while True:
        try:
            data = next(ds_iterator)
        except StopIteration:
            print("End of iteration")
            break

        print(data)
        batch_size *= 2
        num_samples_seen += data.shape[0]
        ds_iterator = iter(get_dataset(batch_size, num_samples_seen))
        print("New batch size:", batch_size)

if __name__ == "__main__":
    main()

如您所见，您需要通过调用get_dataset来实例化一个新的数据集并更新迭代器。

我不知道这种解决方案的性能影响。也许有另一种解决方案，只需要实例化一个batch步骤而不是整个数据集。

- AlexisBRENON

---

由于这将重新初始化数据集，对我来说这不可行，因为性能会受到影响。 - Himaprasoon

是的，数据集没有快速前进的智能方式，因此这将具有二次运行时间，因为它需要重新制作n个结果以生成结果n+1。 - mdaoust

Answer 3

显然，如果您正在使用.from_generator，您可以在其中手动批处理，但这并没有真正回答您的问题。

我能想到的两种最简单的方法是将批量大小作为数据集的组成部分，并构建所需大小的批次：

import tensorflow as tf

batch_sizes = tf.data.Dataset.range(4)
ds = batch_sizes.map(lambda n: tf.random.normal(shape=[n,3]))

for item in ds:
  print(item.shape)
  print()

(0,3)
(1,3)
(2,3)
(3,3)

或者，基于@PG-N的解决方案，如果您需要一个完全在tf.function内运行的版本，您可以使用tf.TensorArray将它们打包：

import tensorflow as tf 

 class Batcher(tf.Module):
   def __init__(self, ds, batch_size=0):   
     self.it = iter(ds) 
     self._batch_size = tf.Variable(batch_size) 

   @property 
   def batch_size(self): 
     return self._batch_size

   @batch_size.setter  
   def batch_size(self, new_size): 
     self._batch_size.assign(new_size) 

   @tf.function 
   def __call__(self): 
     examples =tf.TensorArray(dtype=tf.int64, size=self.batch_size) 
     for i in tf.range(self.batch_size): 
       examples = examples.write(i, next(self.it)) 

     return examples.stack() 

ds = tf.data.Dataset.range(100)
B = Batcher(ds)
B.batch_size = 5
B().numpy()

array([0, 1, 2, 3, 4])

B.batch_size = 10
B().numpy()

array([ 5,  6,  7,  8,  9, 10, 11, 12, 13, 14])

B.batch_size = 3
B().numpy()

array([ 15,  16,  17])

你可以使用tf.nest在中间做些处理，使其适用于具有多个张量组件的数据集。

此外，根据你的用例，像group_by_window和bucket_by_sequence_length这样的方法可能会有所帮助。它们可以进行一些多尺寸批处理，这可能正是你所需要的，或者实现方式可能为你解决问题提供线索。

- mdaoust