def reduce_func(unused_key, window: tf.data.Dataset): return window.padded_batch( batch_size, padded_shapes=(SourceData( key=tf.TensorShape([]), waveform=tf.TensorShape([None, 1]), waveform_length=tf.TensorShape([]), speaker_id=tf.TensorShape([]), age=tf.TensorShape([]), gender=tf.TensorShape([]), ), SourceData( key=tf.TensorShape([]), waveform=tf.TensorShape([None, 1]), waveform_length=tf.TensorShape([]), speaker_id=tf.TensorShape([]), age=tf.TensorShape([]), gender=tf.TensorShape([]), )), padding_values=(SourceData( key="", waveform=tf.to_float(0), waveform_length=tf.to_int64(0), speaker_id=tf.to_int64(0), age=tf.to_int64(0), gender=tf.to_int64(0), ), SourceData( key="", waveform=tf.to_float(0), waveform_length=tf.to_int64(0), speaker_id=tf.to_int64(0), age=tf.to_int64(0), gender=tf.to_int64(0), )))
def reduce_func(unused_key, window: tf.data.Dataset): return window.padded_batch( batch_size, padded_shapes=(SourceData( id=tf.TensorShape([]), key=tf.TensorShape([]), mel=tf.TensorShape([None, self.hparams.num_mels]), mel_length=tf.TensorShape([]), mel_width=tf.TensorShape([]), text=tf.TensorShape([]), ), TargetData( id=tf.TensorShape([]), key=tf.TensorShape([]), waveform=tf.TensorShape([None, 1]), waveform_length=tf.TensorShape([]), )), padding_values=(SourceData( id=tf.to_int64(0), key="", mel=tf.to_float(0), mel_length=tf.to_int64(0), mel_width=tf.to_int64(0), text="", ), TargetData( id=tf.to_int64(0), key="", waveform=tf.to_float(0), waveform_length=tf.to_int64(0), )))
def process(self, dataset: tf.data.Dataset, batch_size: int): dataset = dataset.map(self.parse, num_parallel_calls=AUTOTUNE) if self.cache: dataset = dataset.cache() if self.shuffle: dataset = dataset.shuffle(self.buffer_size, reshuffle_each_iteration=True) # PADDED BATCH the dataset dataset = dataset.padded_batch( batch_size=batch_size, padded_shapes=( tf.TensorShape([]), tf.TensorShape(self.speech_featurizer.shape), tf.TensorShape([]), tf.TensorShape([None]), tf.TensorShape([]), tf.TensorShape([None]), tf.TensorShape([]), ), padding_values=("", 0., 0, self.text_featurizer.blank, 0, self.text_featurizer.blank, 0), drop_remainder=self.drop_remainder) # PREFETCH to improve speed of input length dataset = dataset.prefetch(AUTOTUNE) self.total_steps = get_num_batches(self.total_steps, batch_size) return dataset
def batch_and_split(dataset: tf.data.Dataset, max_sequence_length: int, batch_size: int) -> tf.data.Dataset: return dataset.padded_batch( batch_size, padded_shapes=[max_sequence_length + 1 ]).map(split_input_target, num_parallel_calls=tf.data.experimental.AUTOTUNE)
def get_bucket_iter(self, dataset: tf.data.Dataset, batch_size=32, train=True) -> tf.data.Dataset: padded_shapes = self._padded_shapes() padding_values = self._padding_values() if train: bucket_boundaries = self._bucket_boundaries(batch_size) bucket_batch_sizes = [batch_size] * (len(bucket_boundaries) + 1) dataset = dataset.apply( tf.data.experimental.bucket_by_sequence_length( self.element_length_func, bucket_boundaries, bucket_batch_sizes, padded_shapes=padded_shapes, padding_values=padding_values)) dataset = dataset.shuffle(100) else: dataset = dataset.padded_batch(batch_size, padded_shapes=padded_shapes) dataset = dataset.map(self._collate_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE) dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE) return dataset
def batch_dataset(dataset: tf.data.Dataset, feature_type2name, feature_name2num, batch_size): """Performs batching on ranking dataset When there's no sparse features, padded_batch() is enough. When there are sparse features, we use batching function specific for sparse features """ padded_shapes, padded_values = _get_padded_shapes_and_values( feature_type2name, feature_name2num) # Use padded_batch() if no sparse features if InputFtrType.SPARSE_FTRS_COLUMN_NAMES not in feature_type2name and InputFtrType.SHALLOW_TOWER_SPARSE_FTRS_COLUMN_NAMES not in feature_type2name: # drop_remainder=True to avoid input batch_size=0 issue in evaluation mode in multi gpu training return dataset.padded_batch(batch_size, padded_shapes=padded_shapes, padding_values=padded_values, drop_remainder=True) sparse_batch_fn = partial(_sparse_batch_fn, feature_type2name=feature_type2name, padded_shapes=padded_shapes, padded_values=padded_values, batch_size=batch_size) # drop_remainder=True to avoid input batch_size=0 issue in evaluation mode in multi gpu training return dataset.window(batch_size, drop_remainder=True).flat_map(sparse_batch_fn)
def batch_dataset(dataset: tf.data.Dataset, model: LineRecognizer, batch_size=32, bucket_boundaries=None, padded=True): # add image widths and text length dataset = dataset.map(lambda i, t: (i, tf.shape(i)[ 1], t, tf.strings.length(t, unit='UTF8_CHAR'))) dataset = dataset.map(lambda image, width, text, length: (image, width, model.encoder.encode(text), length)) output_shapes = (model.image_shape, [], [None], []) if bucket_boundaries: if isinstance(batch_size, int): batch_size = [batch_size] * (len(bucket_boundaries) + 1) dataset = dataset.apply( tf.data.experimental.bucket_by_sequence_length( lambda i, w, label, length: w, bucket_boundaries=bucket_boundaries, bucket_batch_sizes=batch_size, padded_shapes=output_shapes)) elif padded: dataset = dataset.padded_batch(batch_size=batch_size, padded_shapes=output_shapes) else: dataset = dataset.batch(batch_size) return dataset
def run_evaluate(self, data: tf.data.Dataset): data = data.padded_batch(self.args.batch) metrics = defaultdict(list) for batch_id, batch in enumerate(data): out = self(batch, training=True) preds = self.extract_preds(out, batch) for k, v in self.compute_metrics(preds, batch).items(): metrics[k].append(v) return {k: sum(v) / len(v) for k, v in metrics.items()}
def apply(self, dataset: tf.data.Dataset, mode: str = None): # pylint: disable=unused-argument padded_shapes = {field.name: field.shape for field in self.fields} padding_values = { field.name: tf.constant(field.default, field.dtype) for field in self.fields } return dataset.padded_batch( self.batch_size, padded_shapes=padded_shapes, padding_values=padding_values, drop_remainder=self.drop_remainder, )
def _hook_dataset_post_precessing(self, ds: tf.data.Dataset, batch_size: int): ds = ds.padded_batch(batch_size=batch_size, padded_shapes=[None]) ds = ds.prefetch(tf.data.experimental.AUTOTUNE) def add_padding_mask(source): enc_padding_mask = create_padding_mask_fm(source) return (source, enc_padding_mask), source ds = ds.map(map_func=add_padding_mask) # Add pad mask # For masked language model task, all datasets are masked return self._apply_mask_for_mlm(ds=ds, vocab_size=self._vocab_size)
def pad_dataset(data: tf.data.Dataset, bs: int, max_length: int, train: bool): data = data.padded_batch(batch_size=bs, padded_shapes=({ "input_ids": [max_length], "segment_ids": [max_length], "input_mask": [max_length] }, []), padding_values=({ "input_ids": 0, "segment_ids": 0, "input_mask": 0, }, 0), drop_remainder=True) return data
def get_data_iter(self, dataset: tf.data.Dataset, batch_size=32, train=True) -> tf.data.Dataset: padded_shapes = self._padded_shapes() padding_values = self._padding_values() if train: dataset = dataset.shuffle(batch_size * 100) dataset = dataset.padded_batch(batch_size, padded_shapes=padded_shapes, padding_values=padding_values) dataset = dataset.map(self._collate_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE) dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE) return dataset
def prepare_dataset( dataset: tf.data.Dataset, model_image_size: Tuple[int, int], augmentation_fn: Optional[ImageDataMapFn] = None, num_epochs: Optional[int] = None, batch_size: Optional[int] = None, shuffle_buffer_size: Optional[int] = None, num_parallel_calls: Optional[int] = None, prefetch_buffer_size: Optional[int] = None, prefetch_to_device: Optional[str] = None, ) -> tf.data.Dataset: # apply data augmentation: if augmentation_fn is not None: dataset = dataset.map( map_image_data(augmentation_fn), num_parallel_calls=num_parallel_calls, ) if shuffle_buffer_size is not None: dataset = dataset.shuffle(buffer_size=shuffle_buffer_size) dataset = dataset.repeat(num_epochs) dataset = dataset.map( map_image_data(prepare_for_batching(model_image_size)), num_parallel_calls=num_parallel_calls, ) # batching and padding if batch_size is not None: dataset = dataset.padded_batch( batch_size=batch_size, padded_shapes=get_padding_shapes( dataset, spatial_image_shape=model_image_size), drop_remainder=True, ) # try to prefetch dataset on certain device if prefetch_to_device is not None: prefetch_fn = tf.data.experimental.prefetch_to_device( device=prefetch_to_device, buffer_size=prefetch_buffer_size) dataset = dataset.apply(prefetch_fn) else: if prefetch_buffer_size is not None: dataset = dataset.prefetch(buffer_size=prefetch_buffer_size) return dataset
def preprocess_wmt_data(dataset: tf.data.Dataset, shuffle: bool, num_epochs: Optional[int] = 1, pack_examples: bool = True, shuffle_buffer_size: int = 1024, max_length: int = 512, batch_size: int = 256, drop_remainder: bool = True, prefetch_size: int = AUTOTUNE): """Shuffle and batch/pack the given dataset.""" def length_filter(max_len): def filter_fn(x): source, target = x['inputs'], x['targets'] l = tf.maximum(tf.shape(source)[0], tf.shape(target)[0]) return tf.less(l, max_len + 1) return filter_fn if max_length > 0: dataset = dataset.filter(length_filter(max_length)) if shuffle: dataset = dataset.shuffle(shuffle_buffer_size) dataset = dataset.repeat(num_epochs) if pack_examples: dataset = pack_dataset(dataset, max_length) dataset = dataset.batch(batch_size, drop_remainder=drop_remainder) else: # simple (static-shape) padded batching dataset = dataset.padded_batch(batch_size, padded_shapes={ 'inputs': max_length, 'targets': max_length }, padding_values={ 'inputs': 0, 'targets': 0 }, drop_remainder=drop_remainder) if prefetch_size: dataset = dataset.prefetch(prefetch_size) return dataset
def reduce_func(unused_key, window: tf.data.Dataset): return window.padded_batch(batch_size, padded_shapes=( SourceData( id=tf.TensorShape([]), key=tf.TensorShape([]), source=tf.TensorShape([None]), source_length=tf.TensorShape([]), speaker_id=tf.TensorShape([]), age=tf.TensorShape([]), gender=tf.TensorShape([]), text=tf.TensorShape([]), ), MelData( id=tf.TensorShape([]), key=tf.TensorShape([]), mel=tf.TensorShape([None, self.hparams.num_mels]), mel_width=tf.TensorShape([]), target_length=tf.TensorShape([]), done=tf.TensorShape([None]), spec_loss_mask=tf.TensorShape([None]), binary_loss_mask=tf.TensorShape([None]), )), padding_values=( SourceData( id=tf.to_int64(0), key="", source=tf.to_int64(0), source_length=tf.to_int64(0), speaker_id=tf.to_int64(0), age=tf.to_int64(0), gender=tf.to_int64(-1), text="", ), MelData( id=tf.to_int64(0), key="", mel=tf.to_float(self.hparams.silence_mel_level_db), mel_width=tf.to_int64(0), target_length=tf.to_int64(0), done=tf.to_float(1), spec_loss_mask=tf.to_float(0), binary_loss_mask=tf.to_float(0), )))
def combine_samples_to_batch(self, data: tf.data.Dataset, batch_size: int) -> tf.data.Dataset: """ Apply batching on data Parameters ---------- data data batch_size batch size Returns ------- data_batch batched data """ data = data.padded_batch(batch_size, data.output_shapes, drop_remainder=self.fix_batch_dimension) return data
def preprocess_wmt_data(dataset: tf.data.Dataset, data_rng, train: bool, shuffle_buffer_size: int = 1024, max_length: int = 512, per_device_batch_size: int = 256): """Shuffle and batch/pack the given dataset.""" def length_filter(max_len): def filter_fn(x): source, target = x['inputs'], x['targets'] l = tf.maximum(tf.shape(source)[0], tf.shape(target)[0]) return tf.less(l, max_len + 1) return filter_fn if max_length > 0: dataset = dataset.filter(length_filter(max_length)) if train: dataset = dataset.shuffle(shuffle_buffer_size, seed=data_rng[0]) dataset = dataset.repeat() dataset = pack_dataset(dataset, max_length) dataset = dataset.batch(per_device_batch_size, drop_remainder=train) else: # simple (static-shape) padded batching dataset = dataset.padded_batch(per_device_batch_size, padded_shapes={ 'inputs': max_length, 'targets': max_length }, padding_values={ 'inputs': 0, 'targets': 0 }, drop_remainder=train) dataset = dataset.prefetch(AUTOTUNE) return dataset
def pack_dataset(dataset: tf.data.Dataset, key2length: Union[int, Dict[str, int]], keys: Optional[List[str]] = None) -> tf.data.Dataset: """Creates a 'packed' version of a dataset on-the-fly. Adapted from the mesh-tf implementation. This is meant to replace the irritation of having to create a separate "packed" version of a dataset to train efficiently on TPU. Each example in the output dataset represents several examples in the input dataset. For each key in the input dataset, two additional keys are created: <key>_segmentation: an int32 tensor identifying the parts representing the original example. <key>_position: an int32 tensor identifying the position within the original example. Example: Two input examples get combined to form an output example. The input examples are: {"inputs": [8, 7, 1, 0], "targets":[4, 1, 0]} {"inputs": [2, 3, 4, 1], "targets":[5, 6, 1]} The output example is: { "inputs": [8, 7, 1, 2, 3, 4, 1, 0, 0, 0] "inputs_segmentations": [1, 1, 1, 2, 2, 2, 2, 0, 0, 0] "inputs_positions": [0, 1, 2, 0, 1, 2, 3, 0, 0, 0] "targets": [4, 1, 5, 6, 1, 0, 0, 0, 0, 0] "targets_segmentations": [1, 1, 2, 2, 2, 0, 0, 0, 0, 0] "targets_positions": [0, 1, 0, 1, 2, 0, 0, 0, 0, 0] } 0 represents padding in both the inputs and the outputs. Sequences in the incoming examples are truncated to length "length", and the sequences in the output examples all have fixed (padded) length "length". Args: dataset: a tf.data.Dataset key2length: an integer, or a dict from feature-key to integer keys: a list of strings (e.g. ["inputs", "targets"]) Returns: a tf.data.Dataset """ shapes = tf.nest.map_structure(lambda spec: spec.shape, dataset.element_spec) if keys is None: keys = list(shapes.keys()) for k in keys: if k not in shapes: raise ValueError( f'Key {k} not found in dataset. Available keys are {shapes.keys()}' ) if not shapes[k].is_compatible_with(tf.TensorShape([None])): raise ValueError('Tensors to be packed must be one-dimensional.') # make sure that the length dictionary contains all keys as well as the # keys suffixed by "_segmentation" and "_position" if isinstance(key2length, int): key2length = {k: key2length for k in keys} for k in keys: for suffix in ['_segmentation', '_position']: key2length[k + suffix] = key2length[k] # trim to length dataset = dataset.map(lambda x: {k: x[k][:key2length[k]] for k in keys}, num_parallel_calls=AUTOTUNE) # Setting batch_size=length ensures that the concatenated sequences (if they # have length >=1) are sufficient to fill at least one packed example. batch_size = max(key2length.values()) dataset = dataset.padded_batch(batch_size, padded_shapes={k: [-1] for k in keys}) dataset = _pack_with_tf_ops(dataset, keys, key2length) # Set the Tensor shapes correctly since they get lost in the process. def my_fn(x): return {k: tf.reshape(v, [key2length[k]]) for k, v in x.items()} return dataset.map(my_fn, num_parallel_calls=AUTOTUNE)
def dataset_batch(ds_single: tf.data.Dataset, batch_shape, batch_size=5): ds = ds_single.padded_batch(batch_size=batch_size, padded_shapes=batch_shape, drop_remainder=False) return ds
def __batch_and_prefetch(self, dataset: tf.data.Dataset) -> tf.data.Dataset: return dataset.padded_batch( self.batch_size, padded_shapes=([None, None, 3], [self.no_classes, ])).prefetch( buffer_size=self.buffer_size)