class_docstr: Selectively backpropagate gradients from a subset of each batch.
  
  Based on (`Jiang et al, 2019`_), Selective Backprop (SB) prunes minibatches
  according to the difficulty of the individual training examples, and only
  computes weight gradients over the pruned subset, reducing iteration time, and
  speeding up training.
  
  The fraction of the minibatch that is kept for gradient computation is
  specified by the argument ``0 <= keep <= 1``.
  
  To speed up SB's selection forward pass, the argument ``scale_factor`` can
  be used to spatially downsample input image tensors. The full-sized inputs
  will still be used for the weight gradient computation.
  
  To preserve convergence, SB can be interrupted with vanilla minibatch
  gradient steps every ``interrupt`` steps. When ``interrupt=0``, SB will be
  used at every step during the SB interval. When ``interrupt=2``, SB will
  alternate with vanilla minibatch steps.
  
  .. _Jiang et al, 2019: https://arxiv.org/abs/1910.00762
  
  Args:
  ----
      start (float, optional): SB interval start as fraction of training duration.
          Default: ``0.5``.
      end (float, optional): SB interval end as fraction of training duration.
          Default: ``0.9``.
      keep (float, optional): fraction of minibatch to select and keep for gradient computation.
          Default: ``0.5``.
      scale_factor (float, optional): scale for downsampling input for selection forward pass.
          Default: ``1.``.
      interrupt (int, optional): interrupt SB with a vanilla minibatch step every
          ``interrupt`` batches. Default: ``2``.
      loss_fn (Callable): Loss function of the form
          ``loss(outputs, targets, reduction='none')``.
          The function must take the keyword argument ``reduction='none'``
          to ensure that per-sample losses are returned.
  
  Example:
  -------
      .. testcode::
  
          from bigdl.nano.pytorch.algorithms.selective_backprop import SelectiveBackprop
          from bigdl.nano.pytorch import Trainer
          from torch import nn
          loss_fn = nn.CrossEntropyLoss(reduction='none')
          sb = SelectiveBackprop(start=0.5, end=0.9, keep=0.5, loss_fn=loss_fn)
          trainer = Trainer(
              algorithms=[sb],
          )
class_docstr_tok: ["Selectively","backpropagate","gradients","from","a","subset","of","each","batch",".","Based","on","(","`","Jiang","et","al",",","2019","`","_",")",",","Selective","Backprop","(","SB",")","prunes","minibatches","according","to","the","difficulty","of","the","individual","training","examples",",","and","only","computes","weight","gradients","over","the","pruned","subset",",","reducing","iteration","time",",","and","speeding","up","training",".","The","fraction","of","the","minibatch","that","is","kept","for","gradient","computation","is","specified","by","the","argument","``","0","<","=","keep","<","=","1","``",".","To","speed","up","SB","'s","selection","forward","pass",",","the","argument","``","scale_factor","``","can","be","used","to","spatially","downsample","input","image","tensors",".","The","full-sized","inputs","will","still","be","used","for","the","weight","gradient","computation",".","To","preserve","convergence",",","SB","can","be","interrupted","with","vanilla","minibatch","gradient","steps","every","``","interrupt","``","steps",".","When","``","interrupt=0","``",",","SB","will","be","used","at","every","step","during","the","SB","interval",".","When","``","interrupt=2","``",",","SB","will","alternate","with","vanilla","minibatch","steps",".","..","_Jiang","et","al",",","2019",":","https",":","\/\/arxiv.org\/abs\/1910.00762","Args",":","--","--","start","(","float",",","optional",")",":","SB","interval","start","as","fraction","of","training","duration",".","Default",":","``","0.5","``",".","end","(","float",",","optional",")",":","SB","interval","end","as","fraction","of","training","duration",".","Default",":","``","0.9","``",".","keep","(","float",",","optional",")",":","fraction","of","minibatch","to","select","and","keep","for","gradient","computation",".","Default",":","``","0.5","``",".","scale_factor","(","float",",","optional",")",":","scale","for","downsampling","input","for","selection","forward","pass",".","Default",":","``","1.","``",".","interrupt","(","int",",","optional",")",":","interrupt","SB","with","a","vanilla","minibatch","step","every","``","interrupt","``","batches",".","Default",":","``","2","``",".","loss_fn","(","Callable",")",":","Loss","function","of","the","form","``","loss","(","outputs",",","targets",",","reduction='none","'",")","``",".","The","function","must","take","the","keyword","argument","``","reduction='none","'","``","to","ensure","that","per-sample","losses","are","returned",".","Example",":","--","--","--","-","..","testcode",":",":","from","bigdl.nano.pytorch.algorithms.selective_backprop","import","SelectiveBackprop","from","bigdl.nano.pytorch","import","Trainer","from","torch","import","nn","loss_fn","=","nn.CrossEntropyLoss","(","reduction='none","'",")","sb","=","SelectiveBackprop","(","start=0.5",",","end=0.9",",","keep=0.5",",","loss_fn=loss_fn",")","trainer","=","Trainer","(","algorithms=","[","sb","]",",",")"]
docstr: Add PyTorch Lightning callback.
docstr_tok: ["Add","PyTorch","Lightning","callback","."]
code: def on_train_batch_start(
      self,
      trainer: "pl.Trainer",
      pl_module: "pl.LightningModule",
      batch: Union[torch.Tensor, torch.Tensor],
      batch_idx: int,
      unused: Any = 0,
  ):
      """Add PyTorch Lightning callback."""
      if self.__match(trainer, batch_idx):
          input, target = batch[0], batch[1]
          if not isinstance(input, torch.Tensor) and isinstance(
              target, torch.Tensor
          ):
              invalidInputError(
                  False,
                  "Multiple tensors \
                      not supported for this method yet.",
              )
  
          input, target = select_using_loss(
              batch,
              batch_idx,
              trainer,
              pl_module,
              self.keep,
              self.scale_factor,
              self._loss_fn,
          )
          batch[0] = input
          batch[1] = target
  
file_path: BigDL/python/nano/src/bigdl/nano/pytorch/algorithms/selective_backprop/selective_backprop.py
imports: from typing import Any, Callable, Tuple, Union
  import warnings
  import numpy
  from bigdl.nano.utils.common import invalidInputError
  import torch
  from torch.nn import functional
  import pytorch_lightning
  from pytorch_lightning.callbacks import Callback
label_desc: Use image node_id 3 for calling the SelectiveBackprop obj's underlying member method code with example usage: obj.on_train_batch_start(trainer, pl_module, batch, batch_idx, unused) without return types