Skip to content

Commit 8133563

Browse files
authored
[IO] Add hb.data.rebatch for resizing batches from tabular data.
This patch introduces `hb.data.rebatch` for resizing batches from tabular data. **Example**: ```python import tensorflow as tf import hybridbackend.tensorflow as hb ds = hb.data.ParquetDataset(filename, micro_batch_size) ds = ds.apply(hb.data.rebatch(batch_size)) ```
1 parent d4ab633 commit 8133563

10 files changed

Lines changed: 1252 additions & 2 deletions

File tree

docs/tutorials/data_loading.md

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ Text | `string`
3636
:special-members: __init__
3737
.. autofunction:: hybridbackend.tensorflow.data.read_parquet
3838
.. autofunction:: hybridbackend.tensorflow.data.to_sparse
39+
.. autofunction:: hybridbackend.tensorflow.data.rebatch
3940
```
4041

4142
## Data Loading Use Cases
@@ -89,6 +90,26 @@ batch = it.get_next()
8990
...
9091
```
9192

93+
### Read and shuffle samples
94+
95+
```python
96+
import tensorflow as tf
97+
import hybridbackend.tensorflow as hb
98+
filenames = tf.data.Dataset.from_generator(func, tf.string, tf.TensorShape([]))
99+
fields = [
100+
hb.data.DataFrame.Field('A', tf.int64),
101+
hb.data.DataFrame.Field('C', tf.int64, ragged_rank=1)]
102+
ds = filenames.apply(hb.data.read_parquet(1024, fields=fields))
103+
ds = ds.shuffle(2048 // 256)
104+
ds = ds.apply(hb.data.rebatch(1024, fields=fields))
105+
ds = ds.apply(hb.data.to_sparse())
106+
ds = ds.prefetch(4)
107+
it = tf.data.make_one_shot_iterator(ds)
108+
batch = it.get_next()
109+
# {'a': tensora, 'c': tensorc}
110+
...
111+
```
112+
92113
## Performance
93114

94115
In benchmark for reading 20k samples from 200 columns of a Parquet file,

0 commit comments

Comments
 (0)