class_docstr: In Appendix F, Page 14 of the TracIn paper, they show that the calculation
  of the influence score of between a test example x' and a training example x,
  can be computed much more quickly than naive back-propagation in the special
  case when considering only gradients in the last fully-connected layer. This class
  computes influence scores for that special case. Note that the computed
  influence scores are exactly the same as when naive back-propagation is used -
  there is no loss in accuracy.
  
  In more detail regarding the influence score computation: let :math`x`
  and :math`\nabla_y f(y)` be the input and output-gradient of the last
  fully-connected layer, respectively, for a training example. Similarly, let
  :math`x'` and :math`\nabla_{y'} f(y')` be the corresponding quantities for
  a test example. Then, the influence score of the training example on the test
  example is the sum of the contribution from each checkpoint. The contribution from
  a given checkpoint is :math`(x^T x')(\nabla_y f(y)^T \nabla_{y'} f(y'))`.
class_docstr_tok: ["In","Appendix","F",",","Page","14","of","the","TracIn","paper",",","they","show","that","the","calculation","of","the","influence","score","of","between","a","test","example","x","'","and","a","training","example","x",",","can","be","computed","much","more","quickly","than","naive","back-propagation","in","the","special","case","when","considering","only","gradients","in","the","last","fully-connected","layer",".","This","class","computes","influence","scores","for","that","special","case",".","Note","that","the","computed","influence","scores","are","exactly","the","same","as","when","naive","back-propagation","is","used","-","there","is","no","loss","in","accuracy",".","In","more","detail","regarding","the","influence","score","computation",":","let",":","math","`","x","`","and",":","math","`","\\nabla_y","f","(","y",")","`","be","the","input","and","output-gradient","of","the","last","fully-connected","layer",",","respectively",",","for","a","training","example",".","Similarly",",","let",":","math","`","x","'","`","and",":","math","`","\\nabla_","{","y","'","}","f","(","y","'",")","`","be","the","corresponding","quantities","for","a","test","example",".","Then",",","the","influence","score","of","the","training","example","on","the","test","example","is","the","sum","of","the","contribution","from","each","checkpoint",".","The","contribution","from","a","given","checkpoint","is",":","math","`","(","x^T","x","'",")","(","\\nabla_y","f","(","y",")","^T","\\nabla_","{","y","'","}","f","(","y","'",")",")","`","."]
docstr: Args:
  
      model (torch.nn.Module): An instance of pytorch model. This model should
              define all of its layers as attributes of the model.
      final_fc_layer (torch.nn.Module or str): The last fully connected layer in
              the network for which gradients will be approximated via fast random
              projection method. Can be either the layer module itself, or the
              fully qualified name of the layer if it is a defined attribute of
              the passed `model`.
      train_dataset (torch.utils.data.Dataset or torch.utils.data.DataLoader):
              In the `influence` method, we compute the influence score of
              training examples on examples in a test batch.
              This argument represents the training dataset containing those
              training examples. In order to compute those influence scores, we
              will create a Pytorch DataLoader yielding batches of training
              examples that is then used for processing. If this argument is
              already a Pytorch Dataloader, that DataLoader can be directly
              used for processing. If it is instead a Pytorch Dataset, we will
              create a DataLoader using it, with batch size specified by
              `batch_size`. For efficiency purposes, the batch size of the
              DataLoader used for processing should be as large as possible, but
              not too large, so that certain intermediate quantities created
              from a batch still fit in memory. Therefore, if
              `train_dataset` is a Dataset, `batch_size` should be large.
              If `train_dataset` was already a DataLoader to begin with,
              it should have been constructed to have a large batch size. It is
              assumed that the Dataloader (regardless of whether it is created
              from a Pytorch Dataset or not) yields tuples. For a `batch` that is
              yielded, of length `L`, it is assumed that the forward function of
              `model` accepts `L-1` arguments, and the last element of `batch` is
              the label. In other words, `model(*batch[:-1])` gives the output of
              `model`, and `batch[-1]` are the labels for the batch.
      checkpoints (str, list[str], or Iterator): Either the directory of the
              path to store and retrieve model checkpoints, a list of
              filepaths with checkpoints from which to load, or an iterator which
              returns objects from which to load checkpoints.
      checkpoints_load_func (Callable, optional): The function to load a saved
              checkpoint into a model to update its parameters, and get the
              learning rate if it is saved. By default uses a utility to load a
              model saved as a state dict.
              Default: _load_flexible_state_dict
      loss_fn (Callable, optional): The loss function applied to model. `loss_fn`
              must be a "reduction" loss function that reduces the per-example
              losses in a batch, and returns a single scalar Tensor. Furthermore,
              the reduction must be the *sum* or the *mean* of the per-example
              losses. For instance, `nn.BCELoss(reduction="sum")` is acceptable.
              Also note that if `loss_fn` has no "reduction" attribute,
              the implementation assumes that the reduction is the *sum* of the
              per-example losses.  If this is not the case, i.e. the reduction
              is the *mean*, please set the "reduction" attribute of `loss_fn`
              to "mean", i.e. `loss_fn.reduction = "mean"`.
              Default: None
      batch_size (int or None, optional): Batch size of the DataLoader created to
              iterate through `train_dataset`, if it is a Dataset.
              `batch_size` should be chosen as large as possible so that certain
              intermediate quantities created from a batch still fit in memory.
              Specific implementations of `TracInCPBase` will detail the size of
              the intermediate quantities. `batch_size` must be an int if
              `train_dataset` is a Dataset. If `train_dataset`
              is a DataLoader, then `batch_size` is ignored as an argument.
              Default: 1
      test_loss_fn (Callable, optional): In some cases, one may want to use a
              separate loss functions for training examples, i.e. those in
              `train_dataset`, and for test examples, i.e. those
              represented by the `inputs` and `targets` arguments to the
              `influence` method. For example, if one wants to calculate the
              influence score of a training example on a test example's
              prediction for a fixed class, `test_loss_fn` could map from the
              logits for all classes to the logits for a fixed class.
              `test_loss_fn` needs satisfy the same constraints as `loss_fn`.
              Thus, the same checks that we apply to `loss_fn` are also applied
              to `test_loss_fn`, if the latter is provided. If not provided, the
              loss function for test examples is assumed to be the same as the
              loss function for training examples, i.e. `loss_fn`.
              Default: None
      vectorize (bool, optional): Flag to use experimental vectorize functionality
              for `torch.autograd.functional.jacobian`.
              Default: False
docstr_tok: ["Args",":","model","(","torch.nn.Module",")",":","An","instance","of","pytorch","model",".","This","model","should","define","all","of","its","layers","as","attributes","of","the","model",".","final_fc_layer","(","torch.nn.Module","or","str",")",":","The","last","fully","connected","layer","in","the","network","for","which","gradients","will","be","approximated","via","fast","random","projection","method",".","Can","be","either","the","layer","module","itself",",","or","the","fully","qualified","name","of","the","layer","if","it","is","a","defined","attribute","of","the","passed","`","model","`",".","train_dataset","(","torch.utils.data.Dataset","or","torch.utils.data.DataLoader",")",":","In","the","`","influence","`","method",",","we","compute","the","influence","score","of","training","examples","on","examples","in","a","test","batch",".","This","argument","represents","the","training","dataset","containing","those","training","examples",".","In","order","to","compute","those","influence","scores",",","we","will","create","a","Pytorch","DataLoader","yielding","batches","of","training","examples","that","is","then","used","for","processing",".","If","this","argument","is","already","a","Pytorch","Dataloader",",","that","DataLoader","can","be","directly","used","for","processing",".","If","it","is","instead","a","Pytorch","Dataset",",","we","will","create","a","DataLoader","using","it",",","with","batch","size","specified","by","`","batch_size","`",".","For","efficiency","purposes",",","the","batch","size","of","the","DataLoader","used","for","processing","should","be","as","large","as","possible",",","but","not","too","large",",","so","that","certain","intermediate","quantities","created","from","a","batch","still","fit","in","memory",".","Therefore",",","if","`","train_dataset","`","is","a","Dataset",",","`","batch_size","`","should","be","large",".","If","`","train_dataset","`","was","already","a","DataLoader","to","begin","with",",","it","should","have","been","constructed","to","have","a","large","batch","size",".","It","is","assumed","that","the","Dataloader","(","regardless","of","whether","it","is","created","from","a","Pytorch","Dataset","or","not",")","yields","tuples",".","For","a","`","batch","`","that","is","yielded",",","of","length","`","L","`",",","it","is","assumed","that","the","forward","function","of","`","model","`","accepts","`","L-1","`","arguments",",","and","the","last","element","of","`","batch","`","is","the","label",".","In","other","words",",","`","model","(","*","batch","[",":","-1","]",")","`","gives","the","output","of","`","model","`",",","and","`","batch","[","-1","]","`","are","the","labels","for","the","batch",".","checkpoints","(","str",",","list","[","str","]",",","or","Iterator",")",":","Either","the","directory","of","the","path","to","store","and","retrieve","model","checkpoints",",","a","list","of","filepaths","with","checkpoints","from","which","to","load",",","or","an","iterator","which","returns","objects","from","which","to","load","checkpoints",".","checkpoints_load_func","(","Callable",",","optional",")",":","The","function","to","load","a","saved","checkpoint","into","a","model","to","update","its","parameters",",","and","get","the","learning","rate","if","it","is","saved",".","By","default","uses","a","utility","to","load","a","model","saved","as","a","state","dict",".","Default",":","_load_flexible_state_dict","loss_fn","(","Callable",",","optional",")",":","The","loss","function","applied","to","model",".","`","loss_fn","`","must","be","a","``","reduction","''","loss","function","that","reduces","the","per-example","losses","in","a","batch",",","and","returns","a","single","scalar","Tensor",".","Furthermore",",","the","reduction","must","be","the","*","sum","*","or","the","*","mean","*","of","the","per-example","losses",".","For","instance",",","`","nn.BCELoss","(","reduction=","''","sum","''",")","`","is","acceptable",".","Also","note","that","if","`","loss_fn","`","has","no","``","reduction","''","attribute",",","the","implementation","assumes","that","the","reduction","is","the","*","sum","*","of","the","per-example","losses",".","If","this","is","not","the","case",",","i.e",".","the","reduction","is","the","*","mean","*",",","please","set","the","``","reduction","''","attribute","of","`","loss_fn","`","to","``","mean","''",",","i.e",".","`","loss_fn.reduction","=","``","mean","''","`",".","Default",":","None","batch_size","(","int","or","None",",","optional",")",":","Batch","size","of","the","DataLoader","created","to","iterate","through","`","train_dataset","`",",","if","it","is","a","Dataset",".","`","batch_size","`","should","be","chosen","as","large","as","possible","so","that","certain","intermediate","quantities","created","from","a","batch","still","fit","in","memory",".","Specific","implementations","of","`","TracInCPBase","`","will","detail","the","size","of","the","intermediate","quantities",".","`","batch_size","`","must","be","an","int","if","`","train_dataset","`","is","a","Dataset",".","If","`","train_dataset","`","is","a","DataLoader",",","then","`","batch_size","`","is","ignored","as","an","argument",".","Default",":","1","test_loss_fn","(","Callable",",","optional",")",":","In","some","cases",",","one","may","want","to","use","a","separate","loss","functions","for","training","examples",",","i.e",".","those","in","`","train_dataset","`",",","and","for","test","examples",",","i.e",".","those","represented","by","the","`","inputs","`","and","`","targets","`","arguments","to","the","`","influence","`","method",".","For","example",",","if","one","wants","to","calculate","the","influence","score","of","a","training","example","on","a","test","example's","prediction","for","a","fixed","class",",","`","test_loss_fn","`","could","map","from","the","logits","for","all","classes","to","the","logits","for","a","fixed","class",".","`","test_loss_fn","`","needs","satisfy","the","same","constraints","as","`","loss_fn","`",".","Thus",",","the","same","checks","that","we","apply","to","`","loss_fn","`","are","also","applied","to","`","test_loss_fn","`",",","if","the","latter","is","provided",".","If","not","provided",",","the","loss","function","for","test","examples","is","assumed","to","be","the","same","as","the","loss","function","for","training","examples",",","i.e",".","`","loss_fn","`",".","Default",":","None","vectorize","(","bool",",","optional",")",":","Flag","to","use","experimental","vectorize","functionality","for","`","torch.autograd.functional.jacobian","`",".","Default",":","False"]
code: def __init__(
      self,
      model: Module,
      final_fc_layer: Union[Module, str],
      train_dataset: Union[Dataset, DataLoader],
      checkpoints: Union[str, List[str], Iterator],
      checkpoints_load_func: Callable = _load_flexible_state_dict,
      loss_fn: Optional[Union[Module, Callable]] = None,
      batch_size: Union[int, None] = 1,
      test_loss_fn: Optional[Union[Module, Callable]] = None,
      vectorize: bool = False,
  ) -> None:
      r"""
      Args:
  
          model (torch.nn.Module): An instance of pytorch model. This model should
                  define all of its layers as attributes of the model.
          final_fc_layer (torch.nn.Module or str): The last fully connected layer in
                  the network for which gradients will be approximated via fast random
                  projection method. Can be either the layer module itself, or the
                  fully qualified name of the layer if it is a defined attribute of
                  the passed `model`.
          train_dataset (torch.utils.data.Dataset or torch.utils.data.DataLoader):
                  In the `influence` method, we compute the influence score of
                  training examples on examples in a test batch.
                  This argument represents the training dataset containing those
                  training examples. In order to compute those influence scores, we
                  will create a Pytorch DataLoader yielding batches of training
                  examples that is then used for processing. If this argument is
                  already a Pytorch Dataloader, that DataLoader can be directly
                  used for processing. If it is instead a Pytorch Dataset, we will
                  create a DataLoader using it, with batch size specified by
                  `batch_size`. For efficiency purposes, the batch size of the
                  DataLoader used for processing should be as large as possible, but
                  not too large, so that certain intermediate quantities created
                  from a batch still fit in memory. Therefore, if
                  `train_dataset` is a Dataset, `batch_size` should be large.
                  If `train_dataset` was already a DataLoader to begin with,
                  it should have been constructed to have a large batch size. It is
                  assumed that the Dataloader (regardless of whether it is created
                  from a Pytorch Dataset or not) yields tuples. For a `batch` that is
                  yielded, of length `L`, it is assumed that the forward function of
                  `model` accepts `L-1` arguments, and the last element of `batch` is
                  the label. In other words, `model(*batch[:-1])` gives the output of
                  `model`, and `batch[-1]` are the labels for the batch.
          checkpoints (str, list[str], or Iterator): Either the directory of the
                  path to store and retrieve model checkpoints, a list of
                  filepaths with checkpoints from which to load, or an iterator which
                  returns objects from which to load checkpoints.
          checkpoints_load_func (Callable, optional): The function to load a saved
                  checkpoint into a model to update its parameters, and get the
                  learning rate if it is saved. By default uses a utility to load a
                  model saved as a state dict.
                  Default: _load_flexible_state_dict
          loss_fn (Callable, optional): The loss function applied to model. `loss_fn`
                  must be a "reduction" loss function that reduces the per-example
                  losses in a batch, and returns a single scalar Tensor. Furthermore,
                  the reduction must be the *sum* or the *mean* of the per-example
                  losses. For instance, `nn.BCELoss(reduction="sum")` is acceptable.
                  Also note that if `loss_fn` has no "reduction" attribute,
                  the implementation assumes that the reduction is the *sum* of the
                  per-example losses.  If this is not the case, i.e. the reduction
                  is the *mean*, please set the "reduction" attribute of `loss_fn`
                  to "mean", i.e. `loss_fn.reduction = "mean"`.
                  Default: None
          batch_size (int or None, optional): Batch size of the DataLoader created to
                  iterate through `train_dataset`, if it is a Dataset.
                  `batch_size` should be chosen as large as possible so that certain
                  intermediate quantities created from a batch still fit in memory.
                  Specific implementations of `TracInCPBase` will detail the size of
                  the intermediate quantities. `batch_size` must be an int if
                  `train_dataset` is a Dataset. If `train_dataset`
                  is a DataLoader, then `batch_size` is ignored as an argument.
                  Default: 1
          test_loss_fn (Callable, optional): In some cases, one may want to use a
                  separate loss functions for training examples, i.e. those in
                  `train_dataset`, and for test examples, i.e. those
                  represented by the `inputs` and `targets` arguments to the
                  `influence` method. For example, if one wants to calculate the
                  influence score of a training example on a test example's
                  prediction for a fixed class, `test_loss_fn` could map from the
                  logits for all classes to the logits for a fixed class.
                  `test_loss_fn` needs satisfy the same constraints as `loss_fn`.
                  Thus, the same checks that we apply to `loss_fn` are also applied
                  to `test_loss_fn`, if the latter is provided. If not provided, the
                  loss function for test examples is assumed to be the same as the
                  loss function for training examples, i.e. `loss_fn`.
                  Default: None
          vectorize (bool, optional): Flag to use experimental vectorize functionality
                  for `torch.autograd.functional.jacobian`.
                  Default: False
      """
      TracInCPBase.__init__(
          self,
          model,
          train_dataset,
          checkpoints,
          checkpoints_load_func,
          loss_fn,
          batch_size,
          test_loss_fn,
      )
  
      self.vectorize = vectorize
  
      # TODO: restore prior state
      self.final_fc_layer = final_fc_layer
      if isinstance(self.final_fc_layer, str):
          self.final_fc_layer = _get_module_from_name(
              model, self.final_fc_layer
          )
      assert isinstance(self.final_fc_layer, Module)
      for param in self.final_fc_layer.parameters():
          param.requires_grad = True
  
      assert loss_fn is not None, "loss function must not be none"
  
      # check `loss_fn`
      self.reduction_type = _check_loss_fn(self, loss_fn, "loss_fn")
      # check `test_loss_fn` if it was provided
      self.test_reduction_type = (
          self.reduction_type
          if test_loss_fn is None
          else _check_loss_fn(self, test_loss_fn, "test_loss_fn")
      )
  
imports: import threading
  import warnings
  from collections import defaultdict
  from typing import Any, Callable, cast, Dict, Iterator, List, Optional, Tuple, Union
  import torch
  from captum._utils.common import _get_module_from_name, _sort_key_list
  from captum._utils.gradient import _gather_distributed_tensors
  from captum._utils.progress import NullProgress, progress
  from captum.influence._core.tracincp import _influence_route_to_helpers, KMostInfluentialResults, TracInCPBase
  from captum.influence._utils.common import _check_loss_fn, _format_inputs_dataset, _get_k_most_influential_helper, _jacobian_loss_wrt_inputs, _load_flexible_state_dict, _self_influence_by_batches_helper, _tensor_batch_dot
  from captum.influence._utils.nearest_neighbors import AnnoyNearestNeighbors, NearestNeighbors
  from captum.log import log_usage
  from torch import device, Tensor
  from torch.nn import Module
  from torch.utils.data import DataLoader, Dataset
label_desc: Use image node_id 1 to create a new TracInCPFast object from inherited base classes: TracInCPBase with example: obj = TracInCPFast(model, final_fc_layer, train_dataset, checkpoints, checkpoints_load_func, loss_fn, batch_size, test_loss_fn, vectorize)