import functools import math import numpy as np from keras.src import tree from keras.src.trainers.data_adapters import array_slicing from keras.src.trainers.data_adapters import data_adapter_utils from keras.src.trainers.data_adapters.data_adapter import DataAdapter class ArrayDataAdapter(DataAdapter): """Adapter for array-like objects, e.g. TF/JAX Tensors, NumPy arrays.""" def __init__( self, x, y=None, sample_weight=None, batch_size=None, steps=None, shuffle=False, class_weight=None, ): if not can_convert_arrays((x, y, sample_weight)): raise ValueError( "Expected all elements of `x` to be array-like. " f"Received invalid types: x={x}" ) if sample_weight is not None: if class_weight is not None: raise ValueError( "You cannot `class_weight` and `sample_weight` " "at the same time." ) if tree.is_nested(y): if isinstance(sample_weight, (list, tuple, dict)): try: tree.assert_same_structure(y, sample_weight) except ValueError: raise ValueError( "You should provide one `sample_weight` array per " "output in `y`. The two structures did not match:\n" f"- y: {y}\n" f"- sample_weight: {sample_weight}\n" ) else: is_samplewise = len(sample_weight.shape) == 1 or ( len(sample_weight.shape) == 2 and sample_weight.shape[1] == 1 ) if not is_samplewise: raise ValueError( "For a model with multiple outputs, when providing " "a single `sample_weight` array, it should only " "have one scalar score per sample " "(i.e. shape `(num_samples,)`). If you want to use " "non-scalar sample weights, pass a `sample_weight` " "argument with one array per model output." ) # Replicate the same sample_weight array on all outputs. sample_weight = tree.map_structure( lambda _: sample_weight, y ) if class_weight is not None: if tree.is_nested(y): raise ValueError( "`class_weight` is only supported for Models with a single " "output." ) sample_weight = data_adapter_utils.class_weight_to_sample_weights( y, class_weight ) inputs = data_adapter_utils.pack_x_y_sample_weight(x, y, sample_weight) data_adapter_utils.check_data_cardinality(inputs) num_samples = set(i.shape[0] for i in tree.flatten(inputs)).pop() self._num_samples = num_samples self._inputs = inputs # If batch_size is not passed but steps is, calculate from the input # data. Defaults to `32` for backwards compatibility. if not batch_size: batch_size = int(math.ceil(num_samples / steps)) if steps else 32 self._size = int(math.ceil(num_samples / batch_size)) self._batch_size = batch_size self._partial_batch_size = num_samples % batch_size self._shuffle = shuffle def get_numpy_iterator(self): inputs = array_slicing.convert_to_sliceable( self._inputs, target_backend="numpy" ) def slice_and_convert_to_numpy(sliceable, indices=None): x = sliceable[indices] x = sliceable.convert_to_numpy(x) return x return self._get_iterator(slice_and_convert_to_numpy, inputs) def get_tf_dataset(self): from keras.src.utils.module_utils import tensorflow as tf shuffle = self._shuffle batch_size = self._batch_size num_samples = self._num_samples num_full_batches = int(self._num_samples // batch_size) # Vectorized version of shuffle. # This is a performance improvement over using `from_tensor_slices`. # The indices of the data are shuffled and batched, and these indices # are then zipped with the data and used to extract a batch of the data # at each step. The performance improvements here come from: # 1. vectorized batch using gather # 2. parallelized map # 3. pipelined permutation generation # 4. optimized permutation batching # 5. disabled static optimizations indices_dataset = tf.data.Dataset.range(1) def permutation(_): # It turns out to be more performant to make a new set of indices # rather than reusing the same range Tensor. (presumably because of # buffer forwarding.) indices = tf.range(num_samples, dtype=tf.int64) if shuffle and shuffle != "batch": indices = tf.random.shuffle(indices) return indices # We prefetch a single element. Computing large permutations can take # quite a while so we don't want to wait for prefetching over an epoch # boundary to trigger the next permutation. On the other hand, too many # simultaneous shuffles can contend on a hardware level and degrade all # performance. indices_dataset = indices_dataset.map(permutation).prefetch(1) def slice_batch_indices(indices): """Convert a Tensor of indices into a dataset of batched indices. This step can be accomplished in several ways. The most natural is to slice the Tensor in a Dataset map. (With a condition on the upper index to handle the partial batch.) However it turns out that coercing the Tensor into a shape which is divisible by the batch size (and handling the last partial batch separately) allows for a much more favorable memory access pattern and improved performance. Args: indices: Tensor which determines the data order for an entire epoch. Returns: A Dataset of batched indices. """ num_in_full_batch = num_full_batches * batch_size first_k_indices = tf.slice(indices, [0], [num_in_full_batch]) first_k_indices = tf.reshape( first_k_indices, [num_full_batches, batch_size] ) flat_dataset = tf.data.Dataset.from_tensor_slices(first_k_indices) if self._partial_batch_size: index_remainder = tf.data.Dataset.from_tensors( tf.slice( indices, [num_in_full_batch], [self._partial_batch_size] ) ) flat_dataset = flat_dataset.concatenate(index_remainder) return flat_dataset def slice_inputs(indices_dataset, inputs): """Slice inputs into a Dataset of batches. Given a Dataset of batch indices and the unsliced inputs, this step slices the inputs in a parallelized fashion and produces a dataset of input batches. Args: indices_dataset: A Dataset of batched indices. inputs: A python data structure that contains the inputs, targets, and possibly sample weights. Returns: A Dataset of input batches matching the batch indices. """ inputs = array_slicing.convert_to_sliceable( self._inputs, target_backend="tensorflow" ) inputs = tree.lists_to_tuples(inputs) dataset = tf.data.Dataset.zip( (indices_dataset, tf.data.Dataset.from_tensors(inputs).repeat()) ) def grab_batch(i, data): def grab_one(x): if isinstance(x, array_slicing.TensorflowSparseWrapper): return array_slicing.slice_tensorflow_sparse_wrapper( x, i ) if isinstance(x, (list, tuple, dict)): return None if tf.is_tensor(x): return tf.gather(x, i, axis=0) return x return tree.traverse(grab_one, data) dataset = dataset.map( grab_batch, num_parallel_calls=tf.data.AUTOTUNE ) # Default optimizations are disabled to avoid the overhead of # (unnecessary) input pipeline graph serialization & deserialization options = tf.data.Options() options.experimental_optimization.apply_default_optimizations = ( False ) if self._shuffle: options.experimental_external_state_policy = ( tf.data.experimental.ExternalStatePolicy.IGNORE ) dataset = dataset.with_options(options) return dataset indices_dataset = indices_dataset.flat_map(slice_batch_indices) if shuffle == "batch": indices_dataset = indices_dataset.map(tf.random.shuffle) dataset = slice_inputs(indices_dataset, self._inputs) options = tf.data.Options() options.experimental_distribute.auto_shard_policy = ( tf.data.experimental.AutoShardPolicy.DATA ) dataset = dataset.with_options(options) return dataset.prefetch(tf.data.AUTOTUNE) def get_jax_iterator(self): inputs = array_slicing.convert_to_sliceable( self._inputs, target_backend="jax" ) def slice_and_convert_to_jax(sliceable, indices=None): x = sliceable[indices] x = sliceable.convert_to_jax_compatible(x) return x return self._get_iterator(slice_and_convert_to_jax, inputs) def get_torch_dataloader(self): import torch from keras.src.backend.torch.core import convert_to_tensor class ArrayDataset(torch.utils.data.Dataset): def __init__(self, array): self.array = array def __getitems__(self, indices): def slice_and_convert(sliceable): x = sliceable[indices] x = sliceable.convert_to_torch_compatible(x) x = convert_to_tensor(x) return x return tree.map_structure(slice_and_convert, self.array) def __len__(self): return len(self.array[0]) class RandomBatchSampler(torch.utils.data.Sampler): def __init__(self, sampler): self.sampler = sampler def __iter__(self): for batch in self.sampler: yield [batch[i] for i in torch.randperm(len(batch))] def __len__(self): return len(self.sampler) if self._shuffle == "batch": batch_sampler = RandomBatchSampler( torch.utils.data.BatchSampler( range(self._num_samples), batch_size=self._batch_size, drop_last=False, ) ) elif self._shuffle: batch_sampler = torch.utils.data.BatchSampler( torch.utils.data.RandomSampler(range(self._num_samples)), batch_size=self._batch_size, drop_last=False, ) else: batch_sampler = torch.utils.data.BatchSampler( torch.utils.data.SequentialSampler(range(self._num_samples)), batch_size=self._batch_size, drop_last=False, ) # Because ArrayDataset.__getitems__ returns full batches organized in # the expected structure, there is nothing to collate. def no_op_collate(batch): return batch inputs = array_slicing.convert_to_sliceable( self._inputs, target_backend="torch" ) dataset = ArrayDataset(inputs) return torch.utils.data.DataLoader( dataset, batch_sampler=batch_sampler, collate_fn=no_op_collate ) def _get_iterator(self, slice_and_convert_fn, inputs): global_permutation = None if self._shuffle and self._shuffle != "batch": global_permutation = np.random.permutation(self._num_samples) for i in range(self._size): start = i * self._batch_size stop = min((i + 1) * self._batch_size, self._num_samples) if self._shuffle == "batch": indices = np.random.permutation(stop - start) + start elif self._shuffle: indices = global_permutation[start:stop] else: indices = slice(start, stop) slice_indices_and_convert_fn = functools.partial( slice_and_convert_fn, indices=indices ) yield tree.map_structure(slice_indices_and_convert_fn, inputs) @property def num_batches(self): return self._size @property def batch_size(self): return self._batch_size @property def has_partial_batch(self): return self._partial_batch_size > 0 @property def partial_batch_size(self): return self._partial_batch_size or None def can_convert_arrays(arrays): """Check if array like-inputs can be handled by `ArrayDataAdapter` Args: inputs: Structure of `Tensor`s, NumPy arrays, or tensor-like. Returns: `True` if `arrays` can be handled by `ArrayDataAdapter`, `False` otherwise. """ return all( tree.flatten(tree.map_structure(array_slicing.can_slice_array, arrays)) )