class_docstr: The purpose of this wrapper is to provide extra capabilities for HuggingFace Trainer, so that
  it can output several forward pass for samples in prediction time and hence be able to work with
  baal. For a more detailed description of the arguments refer to (
  https://huggingface.co/transformers/v3.0.2/main_classes/trainer.html)
  
  Args:
      model (transformers.PreTrainedModel): The model to train, evaluate or use for predictions.
      replicate_in_memory: If True, will perform MC-Dropout in a single forward pass.
          It is faster, but more memory expensive. Default: True.
      data_collator (Optional(Callable)): The function to use to from a batch.
      train_dataset (Optional(torch.utils.data.Dataset)): The dataset to use for training.
      eval_dataset (Optional(torch.utils.data.Dataset)): The dataset to use for evaluation.
      tokenizer (Optional(transformers.PreTrainedTokenizer)): a tokenizer provided by huggingface.
      model_init (Optional(Dict)): Model initial weights for fine tuning.
      compute_metrics (Optional(Callable[[EvalPrediction], Dict])): The function that will be
          used to compute metrics at evaluation.
      callbacks (Optional(List[transformers.TrainerCallback])): A list of callbacks to customize
          the training loop.
      optimizers (Optional(Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR])):
          A tuple containing the optimizer and the scheduler to use.
class_docstr_tok: ["The","purpose","of","this","wrapper","is","to","provide","extra","capabilities","for","HuggingFace","Trainer",",","so","that","it","can","output","several","forward","pass","for","samples","in","prediction","time","and","hence","be","able","to","work","with","baal",".","For","a","more","detailed","description","of","the","arguments","refer","to","(","https",":","\/\/huggingface.co\/transformers\/v3.0.2\/main_classes\/trainer.html",")","Args",":","model","(","transformers.PreTrainedModel",")",":","The","model","to","train",",","evaluate","or","use","for","predictions",".","replicate_in_memory",":","If","True",",","will","perform","MC-Dropout","in","a","single","forward","pass",".","It","is","faster",",","but","more","memory","expensive",".","Default",":","True",".","data_collator","(","Optional","(","Callable",")",")",":","The","function","to","use","to","from","a","batch",".","train_dataset","(","Optional","(","torch.utils.data.Dataset",")",")",":","The","dataset","to","use","for","training",".","eval_dataset","(","Optional","(","torch.utils.data.Dataset",")",")",":","The","dataset","to","use","for","evaluation",".","tokenizer","(","Optional","(","transformers.PreTrainedTokenizer",")",")",":","a","tokenizer","provided","by","huggingface",".","model_init","(","Optional","(","Dict",")",")",":","Model","initial","weights","for","fine","tuning",".","compute_metrics","(","Optional","(","Callable","[","[","EvalPrediction","]",",","Dict","]",")",")",":","The","function","that","will","be","used","to","compute","metrics","at","evaluation",".","callbacks","(","Optional","(","List","[","transformers.TrainerCallback","]",")",")",":","A","list","of","callbacks","to","customize","the","training","loop",".","optimizers","(","Optional","(","Tuple","[","torch.optim.Optimizer",",","torch.optim.lr_scheduler.LambdaLR","]",")",")",":","A","tuple","containing","the","optimizer","and","the","scheduler","to","use","."]
docstr: Use the model to predict on a dataset `iterations` time.
  
  Args:
      dataset (Dataset): Dataset to predict on.
      iterations (int): Number of iterations per sample.
      half (bool): If True use half precision.
      ignore_keys (Optional[List[str]]): A list of keys in the output of your model
          (if it is a dictionary) that should be ignored when gathering predictions.
  Notes:
      The "batch" is made of `batch_size` * `iterations` samples.
  
  Returns:
      Generators [batch_size, n_classes, ..., n_iterations].
docstr_tok: ["Use","the","model","to","predict","on","a","dataset","`","iterations","`","time",".","Args",":","dataset","(","Dataset",")",":","Dataset","to","predict","on",".","iterations","(","int",")",":","Number","of","iterations","per","sample",".","half","(","bool",")",":","If","True","use","half","precision",".","ignore_keys","(","Optional","[","List","[","str","]","]",")",":","A","list","of","keys","in","the","output","of","your","model","(","if","it","is","a","dictionary",")","that","should","be","ignored","when","gathering","predictions",".","Notes",":","The","``","batch","''","is","made","of","`","batch_size","`","*","`","iterations","`","samples",".","Returns",":","Generators","[","batch_size",",","n_classes",",","...",",","n_iterations","]","."]
code: def predict_on_dataset_generator(
      self,
      dataset,
      iterations: int = 1,
      half: bool = False,
      ignore_keys: Optional[List[str]] = None,
  ):
      """
      Use the model to predict on a dataset `iterations` time.
  
      Args:
          dataset (Dataset): Dataset to predict on.
          iterations (int): Number of iterations per sample.
          half (bool): If True use half precision.
          ignore_keys (Optional[List[str]]): A list of keys in the output of your model
              (if it is a dictionary) that should be ignored when gathering predictions.
      Notes:
          The "batch" is made of `batch_size` * `iterations` samples.
  
      Returns:
          Generators [batch_size, n_classes, ..., n_iterations].
      """
  
      dataloader = self.get_eval_dataloader(dataset)
      log.info("Start Predict", dataset=len(dataset))
  
      model = self.model
  
      model.eval()
      for step, inputs in enumerate(tqdm(dataloader)):
          if self.replicate_in_memory:
              try:
                  # We perform MC-Dropout in a single pass, fast, but memory intensive.
                  inputs = map_on_tensor(
                      lambda element: map_on_tensor(
                          lambda d: stack_in_memory(d, iterations),
                          element,
                      ),
                      inputs,
                  )
                  _, out, _ = self.prediction_step(
                      model,
                      inputs,
                      prediction_loss_only=False,
                      ignore_keys=ignore_keys,
                  )
              except RuntimeError as e:
                  if "CUDA out of memory" in str(e):
                      raise RuntimeError(
                          """CUDA ran out of memory while Baal tried to replicate data. See the exception above.
                          Use `replicate_in_memory=False` in order to reduce the memory requirements.
                          Note that there will be some speed trade-offs"""
                      ) from e
                  raise e
              out = map_on_tensor(
                  lambda o: o.view([iterations, -1, *o.size()[1:]]), out
              )
          else:
              # We perform a forward pass `iterations` time. Slower, but memory efficient.
              out = [
                  self.prediction_step(
                      model,
                      inputs,
                      prediction_loss_only=False,
                      ignore_keys=ignore_keys,
                  )[1]
                  for _ in range(iterations)
              ]
              out = _stack_preds(out)
          # Swap axes to match Baal [Batch_size, Classes, ..., Iterations]
          out = map_on_tensor(
              lambda o: o.permute(1, *range(3, o.ndimension()), 2, 0),
              out,
          )
          out = map_on_tensor(lambda x: x.detach(), out)
          if half:
              out = map_on_tensor(lambda x: x.half(), out)
          yield map_on_tensor(lambda x: x.cpu().numpy(), out)
  
file_path: baal/baal/transformers_trainer_wrapper.py
imports: from typing import Optional, List, Sequence, Union
  import numpy
  import structlog
  import torch
  from torch import nn
  from tqdm import tqdm
  from transformers import PreTrainedModel, TrainingArguments
  from baal.utils.warnings import raise_warnings_cache_replicated
  from baal.utils.array_utils import stack_in_memory
  from baal.utils.iterutils import map_on_tensor
label_desc: Use image node_id 2 for calling the BaalTransformersTrainer obj's underlying member method code with example usage: obj.predict_on_dataset_generator(dataset, iterations, half, ignore_keys) without return types