A `Dataset` that batches contiguous elements from its input.
| 3252 | |
| 3253 | |
| 3254 | class BatchDataset(UnaryDataset): |
| 3255 | """A `Dataset` that batches contiguous elements from its input.""" |
| 3256 | |
| 3257 | def __init__(self, input_dataset, batch_size, drop_remainder): |
| 3258 | """See `Dataset.batch()` for details.""" |
| 3259 | self._input_dataset = input_dataset |
| 3260 | self._batch_size = ops.convert_to_tensor( |
| 3261 | batch_size, dtype=dtypes.int64, name="batch_size") |
| 3262 | self._drop_remainder = ops.convert_to_tensor( |
| 3263 | drop_remainder, dtype=dtypes.bool, name="drop_remainder") |
| 3264 | |
| 3265 | constant_drop_remainder = tensor_util.constant_value(self._drop_remainder) |
| 3266 | # pylint: disable=protected-access |
| 3267 | if constant_drop_remainder: |
| 3268 | # NOTE(mrry): `constant_drop_remainder` may be `None` (unknown statically) |
| 3269 | # or `False` (explicitly retaining the remainder). |
| 3270 | # pylint: disable=g-long-lambda |
| 3271 | self._structure = nest.map_structure( |
| 3272 | lambda component_spec: component_spec._batch( |
| 3273 | tensor_util.constant_value(self._batch_size)), |
| 3274 | input_dataset.element_spec) |
| 3275 | else: |
| 3276 | self._structure = nest.map_structure( |
| 3277 | lambda component_spec: component_spec._batch(None), |
| 3278 | input_dataset.element_spec) |
| 3279 | variant_tensor = gen_dataset_ops.batch_dataset_v2( |
| 3280 | input_dataset._variant_tensor, |
| 3281 | batch_size=self._batch_size, |
| 3282 | drop_remainder=self._drop_remainder, |
| 3283 | **self._flat_structure) |
| 3284 | super(BatchDataset, self).__init__(input_dataset, variant_tensor) |
| 3285 | |
| 3286 | @property |
| 3287 | def element_spec(self): |
| 3288 | return self._structure |
| 3289 | |
| 3290 | |
| 3291 | class _VariantTracker(tracking.CapturableResource): |