class_docstr: DistilLabelPatcher is a utility class to generate offline distillation dataloader.
  
  Parameters
  ----------
  teacher_predict
      The teacher predict function, given an input and return the logits/predictions/data that distillation needs.
      Usually, it could be an evaluation mode model.
  dataset
      The training dataset, it will be wrapped by a `_UidDataset` in this patcher.
      If this dataset is already an instance of `_UidDataset` and `uid_dataset_cls` is None, it will not be wrap again.
      (Reference to what is _UidDataset ...)
  create_dataloader
      A factory function to create a dataloader for the given `dataset`.
  keep_in_memory
      By default, False. If using memory or disk to save the distillation data (`teacher_predict` outputs).
  cache_folder
      By default, None. The path to save the data if `keep_in_memory` is False. If `cache_folder` is None,
      a temp folder will be created to save data.
      This argument will be ignored if `keep_in_memory` is True.
  cache_mode
      (Not stable!)  By default, `pickle`. Using `pickle` or `hdf5` to save data.
  checkpoint_folder
      By default, None. The path to load storage state to recover the previous saved distillation data.
  uid_dataset_cls
      By default, None. If `dataset` is not an instance of `_UidDataset` and `uid_dataset_cls` is None,
      `IndexedDataset` will be used to wrap `dataset`. `IndexedDataset` take the index passed to `Dataset.__getitem__` as data's uid.
      `uid_dataset_cls` should be a subclass of `_UidDataset`, `uid_dataset_cls.__getitem__` will return a tuple of (uid, original_data).
      (Official uid dataset and how to customize uid dataset, please refer ...)
  uidd_args
      Additional arguments pass to `uid_dataset_cls.__init__`.
      If there are many extra arguments, recommended to directly set `uid_dataset_cls` to None, and then pass in the wrapped `dataset`.
  uidd_kwargs
      Additional keyword arguments pass to `uid_dataset_cls.__init__`.
      If there are many extra arguments, recommended to directly set `uid_dataset_cls` to None, and then pass in the wrapped `dataset`.
  labels_split_fn
      Please refer `DistilLabelPatcher.generate_distillation_labels`.
  labels_collate_fn
      Please refer `DistilLabelPatcher.create_patched_dataloader`.
class_docstr_tok: ["DistilLabelPatcher","is","a","utility","class","to","generate","offline","distillation","dataloader",".","Parameters","--","--","--","--","--","teacher_predict","The","teacher","predict","function",",","given","an","input","and","return","the","logits\/predictions\/data","that","distillation","needs",".","Usually",",","it","could","be","an","evaluation","mode","model",".","dataset","The","training","dataset",",","it","will","be","wrapped","by","a","`","_UidDataset","`","in","this","patcher",".","If","this","dataset","is","already","an","instance","of","`","_UidDataset","`","and","`","uid_dataset_cls","`","is","None",",","it","will","not","be","wrap","again",".","(","Reference","to","what","is","_UidDataset","...",")","create_dataloader","A","factory","function","to","create","a","dataloader","for","the","given","`","dataset","`",".","keep_in_memory","By","default",",","False",".","If","using","memory","or","disk","to","save","the","distillation","data","(","`","teacher_predict","`","outputs",")",".","cache_folder","By","default",",","None",".","The","path","to","save","the","data","if","`","keep_in_memory","`","is","False",".","If","`","cache_folder","`","is","None",",","a","temp","folder","will","be","created","to","save","data",".","This","argument","will","be","ignored","if","`","keep_in_memory","`","is","True",".","cache_mode","(","Not","stable","!",")","By","default",",","`","pickle","`",".","Using","`","pickle","`","or","`","hdf5","`","to","save","data",".","checkpoint_folder","By","default",",","None",".","The","path","to","load","storage","state","to","recover","the","previous","saved","distillation","data",".","uid_dataset_cls","By","default",",","None",".","If","`","dataset","`","is","not","an","instance","of","`","_UidDataset","`","and","`","uid_dataset_cls","`","is","None",",","`","IndexedDataset","`","will","be","used","to","wrap","`","dataset","`",".","`","IndexedDataset","`","take","the","index","passed","to","`","Dataset.__getitem__","`","as","data","'s","uid",".","`","uid_dataset_cls","`","should","be","a","subclass","of","`","_UidDataset","`",",","`","uid_dataset_cls.__getitem__","`","will","return","a","tuple","of","(","uid",",","original_data",")",".","(","Official","uid","dataset","and","how","to","customize","uid","dataset",",","please","refer","...",")","uidd_args","Additional","arguments","pass","to","`","uid_dataset_cls.__init__","`",".","If","there","are","many","extra","arguments",",","recommended","to","directly","set","`","uid_dataset_cls","`","to","None",",","and","then","pass","in","the","wrapped","`","dataset","`",".","uidd_kwargs","Additional","keyword","arguments","pass","to","`","uid_dataset_cls.__init__","`",".","If","there","are","many","extra","arguments",",","recommended","to","directly","set","`","uid_dataset_cls","`","to","None",",","and","then","pass","in","the","wrapped","`","dataset","`",".","labels_split_fn","Please","refer","`","DistilLabelPatcher.generate_distillation_labels","`",".","labels_collate_fn","Please","refer","`","DistilLabelPatcher.create_patched_dataloader","`","."]
docstr: Save a checkpoint that can be used to re-initialize the `DistilLabelPatcher` and reload the generated distillation labels.
  
  Parameters
  ----------
  checkpoint_folder
      The directory to save checkpoint.
docstr_tok: ["Save","a","checkpoint","that","can","be","used","to","re-initialize","the","`","DistilLabelPatcher","`","and","reload","the","generated","distillation","labels",".","Parameters","--","--","--","--","--","checkpoint_folder","The","directory","to","save","checkpoint","."]
code: def save_checkpoint(self, checkpoint_folder: str | PathLike):
      """
      Save a checkpoint that can be used to re-initialize the `DistilLabelPatcher` and reload the generated distillation labels.
  
      Parameters
      ----------
      checkpoint_folder
          The directory to save checkpoint.
      """
      root_path = Path(checkpoint_folder).absolute()
      root_path.mkdir(parents=True, exist_ok=True)
      self._storage.save_checkpoint(root_path)
  
file_path: nni/nni/contrib/distillation/label_patcher.py
imports: from __future__ import annotations
  from collections import abc
  import logging
  from os import PathLike
  from pathlib import Path
  from typing import Any, Callable, Dict, List, Literal, Type, cast
  from torch.utils.data import DataLoader, Dataset
  from .storage import create_storage
  from .uid_dataset import create_uid_dataset, _UidDataset
  from .utils import _default_labels_split_fn, _default_labels_collate_fn
label_desc: Use image node_id 6 for calling the DistilLabelPatcher obj's underlying member method code with example usage: obj.save_checkpoint(checkpoint_folder) without return types