class_docstr: Lightning callback that enables in-training sparsity.
  
  This callback aims to sparsify the model inside lightning module during training.
  **Note that the model is copied and then sparsified, so the existing model is not modified**
  
  The sparsified model can be used for comparison and can be accessed using
      <callback_obj>.sparsified
  
  Args:
      data_sparsifier_class (some implemented class of BaseDataSparsifier)
          The data sparsifier object of this class is created when the
          training starts.
          Note: Objects should not be passed in here as they are created
          when the training starts.
  
      data_sparsifier_args (Dict)
          Dictionary of args to be passed to the data sparsifier.
          Note: data_list arg should be ignored
  
      data_scheduler_class (some implemented class of BaseDataScheduler)
          The data scheduler of this class is created when the training starts
          Note: Objects should not be passed in here as they are created
          when the training starts.
  
      data_scheduler_args(Dict)
          Dictionary of args to be passed to the data scheduler.
          **Note: data_sparsifier arg should be ignored as the recipe
          creates and pass sparsifier object into the class**
  
  Hooks implemented:
      on_train_start()
          Data sparsifier and scheduler objects are created.
          Pytorch model attached to the sparsifier
  
      on_train_epoch_start()
          Loads the state_dict of the data sparsifier
  
      on_train_epoch_end()
          1. Copies the model and attaches it to the sparsifier
          2. sparsifier step() and scheduler step()
          3. Dump state_dict of the current sparsifier
  
      on_train_end()
          squash mask
class_docstr_tok: ["Lightning","callback","that","enables","in-training","sparsity",".","This","callback","aims","to","sparsify","the","model","inside","lightning","module","during","training",".","*","*","Note","that","the","model","is","copied","and","then","sparsified",",","so","the","existing","model","is","not","modified","*","*","The","sparsified","model","can","be","used","for","comparison","and","can","be","accessed","using","<","callback_obj",">",".sparsified","Args",":","data_sparsifier_class","(","some","implemented","class","of","BaseDataSparsifier",")","The","data","sparsifier","object","of","this","class","is","created","when","the","training","starts",".","Note",":","Objects","should","not","be","passed","in","here","as","they","are","created","when","the","training","starts",".","data_sparsifier_args","(","Dict",")","Dictionary","of","args","to","be","passed","to","the","data","sparsifier",".","Note",":","data_list","arg","should","be","ignored","data_scheduler_class","(","some","implemented","class","of","BaseDataScheduler",")","The","data","scheduler","of","this","class","is","created","when","the","training","starts","Note",":","Objects","should","not","be","passed","in","here","as","they","are","created","when","the","training","starts",".","data_scheduler_args","(","Dict",")","Dictionary","of","args","to","be","passed","to","the","data","scheduler",".","*","*","Note",":","data_sparsifier","arg","should","be","ignored","as","the","recipe","creates","and","pass","sparsifier","object","into","the","class","*","*","Hooks","implemented",":","on_train_start","(",")","Data","sparsifier","and","scheduler","objects","are","created",".","Pytorch","model","attached","to","the","sparsifier","on_train_epoch_start","(",")","Loads","the","state_dict","of","the","data","sparsifier","on_train_epoch_end","(",")","1",".","Copies","the","model","and","attaches","it","to","the","sparsifier","2.","sparsifier","step","(",")","and","scheduler","step","(",")","3",".","Dump","state_dict","of","the","current","sparsifier","on_train_end","(",")","squash","mask"]
code: def __init__(
      self,
      data_sparsifier_class,
      data_sparsifier_args,
      data_scheduler_class,
      data_scheduler_args,
  ):
      super().__init__()
      # data sparsifier objects
      self.data_sparsifier_class = data_sparsifier_class
      self.data_sparsifier_args = data_sparsifier_args
  
      # scheduler objects
      self.data_scheduler_class = data_scheduler_class
      self.data_scheduler_args = data_scheduler_args
  
      # fields
      self.data_sparsifier: Any = None
      self.data_scheduler: Any = None
      self.sparsified: Optional[torch.nn.Module] = None
  
      self.data_sparsifier_state_dict: Any = None
  
file_path: pytorch/torch/ao/pruning/_experimental/data_sparsifier/lightning/callbacks/data_sparsity.py
imports: from collections import defaultdict
  from copy import deepcopy
  import torch
  from typing import Any, Optional, Dict
  import pytorch_lightning
  from ._data_sparstity_utils import _attach_model_to_data_sparsifier, _log_sparsified_level, _get_valid_name
label_desc: Use image node_id 1 to create a new TrainingAwareDataSparsity object from inherited base classes: pl.callbacks with example: obj = TrainingAwareDataSparsity(data_sparsifier_class, data_sparsifier_args, data_scheduler_class, data_scheduler_args)