def next_batch(self, batch_size):
    """Gets a batch of items from the input RDD.

    If multiple tensors are provided per row in the input RDD, e.g. tuple of (tensor1, tensor2, ..., tensorN) and:

    * no ``input_mapping`` was provided to the DataFeed constructor, this will return an array of ``batch_size`` tuples,
      and the caller is responsible for separating the tensors.
    * an ``input_mapping`` was provided to the DataFeed constructor, this will return a dictionary of N tensors,
      with tensor names as keys and arrays of length ``batch_size`` as values.

    Note: if the end of the data is reached, this may return with fewer than ``batch_size`` items.

    Args:
      :batch_size: number of items to retrieve.

    Returns:
      A batch of items or a dictionary of tensors.
    """
    logging.debug("next_batch() invoked")
    queue = self.mgr.get_queue(self.qname_in)
    tensors = [] if self.input_tensors is None else {tensor: [] for tensor in self.input_tensors}
    count = 0
    while count < batch_size:
      item = queue.get(block=True)
      if item is None:
        # End of Feed
        logging.info("next_batch() got None")
        queue.task_done()
        self.done_feeding = True
        break
      elif type(item) is marker.EndPartition:
        # End of Partition
        logging.info("next_batch() got EndPartition")
        queue.task_done()
        if not self.train_mode and count > 0:
          break
      else:
        # Normal item
        if self.input_tensors is None:
          tensors.append(item)
        else:
          for i in range(len(self.input_tensors)):
            tensors[self.input_tensors[i]].append(item[i])
        count += 1
        queue.task_done()
    logging.debug("next_batch() returning {0} items".format(count))
    return tensors
def next_batch ( self , batch_size ) : logging . debug ( "next_batch() invoked" ) queue = self . mgr . get_queue ( self . qname_in ) tensors = [ ] if self . input_tensors is None else { tensor : [ ] for tensor in self . input_tensors } count = 0 while count < batch_size : item = queue . get ( block = True ) if item is None : logging . info ( "next_batch() got None" ) queue . task_done ( ) self . done_feeding = True break elif type ( item ) is marker . EndPartition : logging . info ( "next_batch() got EndPartition" ) queue . task_done ( ) if not self . train_mode and count > 0 : break else : if self . input_tensors is None : tensors . append ( item ) else : for i in range ( len ( self . input_tensors ) ) : tensors [ self . input_tensors [ i ] ] . append ( item [ i ] ) count += 1 queue . task_done ( ) logging . debug ( "next_batch() returning {0} items" . format ( count ) ) return tensors
Gets a batch of items from the input RDD.

    If multiple tensors are provided per row in the input RDD, e.g. tuple of (tensor1, tensor2, ..., tensorN) and:

    * no ``input_mapping`` was provided to the DataFeed constructor, this will return an array of ``batch_size`` tuples,
      and the caller is responsible for separating the tensors.
    * an ``input_mapping`` was provided to the DataFeed constructor, this will return a dictionary of N tensors,
      with tensor names as keys and arrays of length ``batch_size`` as values.

    Note: if the end of the data is reached, this may return with fewer than ``batch_size`` items.

    Args:
      :batch_size: number of items to retrieve.

    Returns:
      A batch of items or a dictionary of tensors.