(dataset, training, n_devices, variable_shapes,
             batch_size_per_device=32, batch_size=None, eval_batch_size=32,
             bucket_length=32, buckets=None,
             buckets_include_inputs_in_length=False,
             batch_shuffle_size=None, max_eval_length=None,
             id_to_mask=None, strict_pad_on_len=False)
return add_loss_weights(dataset, id_to_mask)
Batching function.
Batching function.
Batching function .
def batch_fn(dataset, training, n_devices, variable_shapes,
             batch_size_per_device=32, batch_size=None, eval_batch_size=32,
             bucket_length=32, buckets=None,
             buckets_include_inputs_in_length=False,
             batch_shuffle_size=None, max_eval_length=None,
             id_to_mask=None, strict_pad_on_len=False):
  """Batching function."""
  # TODO(lukaszkaiser, jonni): revisit arguments, their semantics and naming.
  # After that, create a proper doc-string; we may also not need to pass both
  # training and eval arguments here, as batcher calls the function separately
  # now and it's not under gin-config any more -- consider reducing args.
  batch_size = batch_size or batch_size_per_device * n_devices
  # If bucketing is not specified, check if target shapes are variable.
  cur_batch_size = batch_size if training else eval_batch_size
  # Make cur_batch_size divisible by n_devices.
  cur_batch_size = max(cur_batch_size // n_devices, 1) * n_devices
  # Create heuristic buckets if none are specified.
  if buckets is None:
    logging.info('Heuristically setting bucketing to %s based on shapes '
                 'of target tensors.', variable_shapes)
    if variable_shapes:
      buckets = _buckets_for_length(
          bucket_length, cur_batch_size, max_eval_length, n_devices, training)

  if buckets:
    logging.info('Bucketing with buckets %s.', str(buckets))
    def example_length(x):
      """The length function used by bucket_by_sequence_length to bucket."""
      # The input x is a tuple to go on the stack, typically either
      # (input, target) or (input, target, mask).
      example_inputs, target = x[0], x[1]
      # Length is the shape of axis 0 here (no batch yet).
      other_length = 0  # We include input length only if asked.
      if buckets_include_inputs_in_length:
        other_length = example_inputs.shape[0]
      return max(target.shape[0], other_length)
    boundaries, batch_sizes = buckets
    dataset = bucket_by_length(
        dataset, example_length, boundaries, batch_sizes, strict_pad_on_len)
  else:
    logging.info('Not Bucketing cur_batch_size %d.', cur_batch_size)
    dataset = batch(dataset, cur_batch_size)
  if training and batch_shuffle_size is not None:
    dataset = shuffle(dataset, batch_shuffle_size)
  return add_loss_weights(dataset, id_to_mask)
def batch_fn ( dataset , training , n_devices , variable_shapes , batch_size_per_device = 32 , batch_size = None , eval_batch_size = 32 , bucket_length = 32 , buckets = None , buckets_include_inputs_in_length = False , batch_shuffle_size = None , max_eval_length = None , id_to_mask = None , strict_pad_on_len = False ) : # TODO(lukaszkaiser, jonni): revisit arguments, their semantics and naming. # After that, create a proper doc-string; we may also not need to pass both # training and eval arguments here, as batcher calls the function separately # now and it's not under gin-config any more -- consider reducing args. batch_size = batch_size or batch_size_per_device * n_devices # If bucketing is not specified, check if target shapes are variable. cur_batch_size = batch_size if training else eval_batch_size # Make cur_batch_size divisible by n_devices. cur_batch_size = max ( cur_batch_size // n_devices , 1 ) * n_devices # Create heuristic buckets if none are specified. if buckets is None : logging . info ( 'Heuristically setting bucketing to %s based on shapes ' 'of target tensors.' , variable_shapes ) if variable_shapes : buckets = _buckets_for_length ( bucket_length , cur_batch_size , max_eval_length , n_devices , training ) if buckets : logging . info ( 'Bucketing with buckets %s.' , str ( buckets ) ) def example_length ( x ) : """The length function used by bucket_by_sequence_length to bucket.""" # The input x is a tuple to go on the stack, typically either # (input, target) or (input, target, mask). example_inputs , target = x [ 0 ] , x [ 1 ] # Length is the shape of axis 0 here (no batch yet). other_length = 0 # We include input length only if asked. if buckets_include_inputs_in_length : other_length = example_inputs . shape [ 0 ] return max ( target . shape [ 0 ] , other_length ) boundaries , batch_sizes = buckets dataset = bucket_by_length ( dataset , example_length , boundaries , batch_sizes , strict_pad_on_len ) else : logging . info ( 'Not Bucketing cur_batch_size %d.' , cur_batch_size ) dataset = batch ( dataset , cur_batch_size ) if training and batch_shuffle_size is not None : dataset = shuffle ( dataset , batch_shuffle_size ) return add_loss_weights ( dataset , id_to_mask )