class_docstr: A palettization algorithm based on `"DKM: Differentiable K-Means Clustering Layer for Neural Network
  Compression" <https://arxiv.org/pdf/2108.12659.pdf>`_. It clusters the weights
  using a differentiable version of ``k-means``, allowing the look-up-table (LUT)
  and indices of palettized weights to be learnt using a gradient-based optimization algorithm such as SGD.
  
  Example:
  
      .. code-block:: python
  
          import torch
          from coremltools.optimize.torch.palettization import (
              DKMPalettizer,
              DKMPalettizerConfig,
              ModuleDKMPalettizerConfig,
          )
  
          # code that defines the pytorch model, loss and optimizer.
          model, loss_fn, optimizer = create_model_loss_and_optimizer()
  
          # initialize the palettizer
          config = DKMPalettizerConfig(global_config=ModuleDKMPalettizerConfig(n_bits=4))
  
          palettizer = DKMPalettizer(model, config)
  
          # prepare the model to insert FakePalettize layers for palettization
          model = palettizer.prepare(inplace=True)
  
          # use palettizer in your PyTorch training loop
          for inputs, labels in data:
              output = model(inputs)
              loss = loss_fn(output, labels)
              loss.backward()
              optimizer.step()
              palettizer.step()
  
          # fold LUT and indices into weights
          model = palettizer.finalize(inplace=True)
  
  Args:
      model (:py:class:`torch.nn.Module`): Model on which the palettizer will act.
      config (:py:class:`DKMPalettizerConfig`): Config which specifies how
          different submodules in the model will be configured for palettization.
          Default config is used when passed as ``None``.
class_docstr_tok: ["A","palettization","algorithm","based","on","`","``","DKM",":","Differentiable","K-Means","Clustering","Layer","for","Neural","Network","Compression","''","<","https",":","\/\/arxiv.org\/pdf\/2108.12659.pdf",">","`","_",".","It","clusters","the","weights","using","a","differentiable","version","of","``","k-means","``",",","allowing","the","look-up-table","(","LUT",")","and","indices","of","palettized","weights","to","be","learnt","using","a","gradient-based","optimization","algorithm","such","as","SGD",".","Example",":","..","code-block",":",":","python","import","torch","from","coremltools.optimize.torch.palettization","import","(","DKMPalettizer",",","DKMPalettizerConfig",",","ModuleDKMPalettizerConfig",",",")","#","code","that","defines","the","pytorch","model",",","loss","and","optimizer",".","model",",","loss_fn",",","optimizer","=","create_model_loss_and_optimizer","(",")","#","initialize","the","palettizer","config","=","DKMPalettizerConfig","(","global_config=ModuleDKMPalettizerConfig","(","n_bits=4",")",")","palettizer","=","DKMPalettizer","(","model",",","config",")","#","prepare","the","model","to","insert","FakePalettize","layers","for","palettization","model","=","palettizer.prepare","(","inplace=True",")","#","use","palettizer","in","your","PyTorch","training","loop","for","inputs",",","labels","in","data",":","output","=","model","(","inputs",")","loss","=","loss_fn","(","output",",","labels",")","loss.backward","(",")","optimizer.step","(",")","palettizer.step","(",")","#","fold","LUT","and","indices","into","weights","model","=","palettizer.finalize","(","inplace=True",")","Args",":","model","(",":","py",":","class",":","`","torch.nn.Module","`",")",":","Model","on","which","the","palettizer","will","act",".","config","(",":","py",":","class",":","`","DKMPalettizerConfig","`",")",":","Config","which","specifies","how","different","submodules","in","the","model","will","be","configured","for","palettization",".","Default","config","is","used","when","passed","as","``","None","``","."]
code: def __init__(
      self,
      model: _nn.Module,
      config: _Optional[_DKMPalettizerConfig] = None,
  ):
      config = _DKMPalettizerConfig() if config is None else config
      super().__init__(model, config)
      self._milestones = {}
      self._supported_modules = _get_palettization_qat_mappings()
  
imports: import copy
  import logging
  from typing import Dict
  from typing import Optional
  import torch
  import torch.nn
  from torch.ao.quantization import FakeQuantize
  from coremltools.optimize.torch._typing import ParamsDict
  from coremltools.optimize.torch._utils.math_utils import rmse_error
  from coremltools.optimize.torch._utils.torch_utils import get_eval_model
  from coremltools.optimize.torch.base_model_optimizer import BaseModelOptimizer
  from coremltools.optimize.torch.base_model_optimizer import _Report
  from coremltools.optimize.torch.palettization._custom_conversion import PALETTIZATION_CONVERT_DICT
  from coremltools.optimize.torch.palettization._supported_modules import _get_palettization_qat_mappings
  from coremltools.optimize.torch.palettization.fake_palettize import FakePalettize
  from coremltools.optimize.torch.palettization.palettization_config import DEFAULT_PALETTIZATION_ADVANCED_OPTIONS
  from coremltools.optimize.torch.palettization.palettization_config import DEFAULT_PALETTIZATION_SCHEME
  from coremltools.optimize.torch.palettization.palettization_config import DKMPalettizerConfig
  from coremltools.optimize.torch.palettization.palettization_config import ModuleDKMPalettizerConfig
label_desc: Use image node_id 1 to create a new DKMPalettizer object from inherited base classes: Palettizer with example: obj = DKMPalettizer(model, config)