class_docstr: Adding a linear layer to a classifier model after the model is
  trained and train this new layer until convergence.
  Together with the linear layer, the model is now calibrated.
  Source: https://arxiv.org/abs/1910.12656
  Code inspired from: https://github.com/dirichletcal/experiments_neurips
  
  References:
      @article{kullbeyond,
              title={Beyond temperature scaling: Obtaining well-calibrated multi-class
               probabilities with Dirichlet calibration Supplementary material},
              author={Kull, Meelis and Perello-Nieto,
               Miquel and K{"a}ngsepp, Markus and Silva Filho,
                Telmo and Song, Hao and Flach, Peter}
              }
  
  Args:
  
      wrapper (ModelWrapper): Provides training and testing methods.
      num_classes (int): Number of classes in classification task.
      lr (float): Learning rate.
      reg_factor (float): Regularization factor for the linear layer weights.
      mu (float): Regularization factor for the linear layer biases.
          If not given, will be initialized by "l".
class_docstr_tok: ["Adding","a","linear","layer","to","a","classifier","model","after","the","model","is","trained","and","train","this","new","layer","until","convergence",".","Together","with","the","linear","layer",",","the","model","is","now","calibrated",".","Source",":","https",":","\/\/arxiv.org\/abs\/1910.12656","Code","inspired","from",":","https",":","\/\/github.com\/dirichletcal\/experiments_neurips","References",":","@","article","{","kullbeyond",",","title=","{","Beyond","temperature","scaling",":","Obtaining","well-calibrated","multi-class","probabilities","with","Dirichlet","calibration","Supplementary","material","}",",","author=","{","Kull",",","Meelis","and","Perello-Nieto",",","Miquel","and","K","{","``","a","}","ngsepp",",","Markus","and","Silva","Filho",",","Telmo","and","Song",",","Hao","and","Flach",",","Peter","}","}","Args",":","wrapper","(","ModelWrapper",")",":","Provides","training","and","testing","methods",".","num_classes","(","int",")",":","Number","of","classes","in","classification","task",".","lr","(","float",")",":","Learning","rate",".","reg_factor","(","float",")",":","Regularization","factor","for","the","linear","layer","weights",".","mu","(","float",")",":","Regularization","factor","for","the","linear","layer","biases",".","If","not","given",",","will","be","initialized","by","``","l","''","."]
docstr: Training the linear layer given a training set and a validation set.
  The training set should be different from what model is trained on.
  
  Args:
      train_set (Dataset): The training set.
      test_set (Dataset): The validation set.
      batch_size (int): Batch size used.
      epoch (int): Number of epochs to train the linear layer for.
      use_cuda (bool): If "True", will use GPU.
      double_fit (bool): If "True" would fit twice on the train set.
      kwargs (dict): Rest of parameters for baal.ModelWrapper.train_and_test_on_dataset().
  
  Returns:
      loss_history (list[float]): List of loss values for each epoch.
      model.state_dict (dict): Model weights.
docstr_tok: ["Training","the","linear","layer","given","a","training","set","and","a","validation","set",".","The","training","set","should","be","different","from","what","model","is","trained","on",".","Args",":","train_set","(","Dataset",")",":","The","training","set",".","test_set","(","Dataset",")",":","The","validation","set",".","batch_size","(","int",")",":","Batch","size","used",".","epoch","(","int",")",":","Number","of","epochs","to","train","the","linear","layer","for",".","use_cuda","(","bool",")",":","If","``","True","''",",","will","use","GPU",".","double_fit","(","bool",")",":","If","``","True","''","would","fit","twice","on","the","train","set",".","kwargs","(","dict",")",":","Rest","of","parameters","for","baal.ModelWrapper.train_and_test_on_dataset","(",")",".","Returns",":","loss_history","(","list","[","float","]",")",":","List","of","loss","values","for","each","epoch",".","model.state_dict","(","dict",")",":","Model","weights","."]
code: def calibrate(
      self,
      train_set: Dataset,
      test_set: Dataset,
      batch_size: int,
      epoch: int,
      use_cuda: bool,
      double_fit: bool = False,
      **kwargs,
  ):
      """
      Training the linear layer given a training set and a validation set.
      The training set should be different from what model is trained on.
  
      Args:
          train_set (Dataset): The training set.
          test_set (Dataset): The validation set.
          batch_size (int): Batch size used.
          epoch (int): Number of epochs to train the linear layer for.
          use_cuda (bool): If "True", will use GPU.
          double_fit (bool): If "True" would fit twice on the train set.
          kwargs (dict): Rest of parameters for baal.ModelWrapper.train_and_test_on_dataset().
  
      Returns:
          loss_history (list[float]): List of loss values for each epoch.
          model.state_dict (dict): Model weights.
  
      """
  
      # reinitialize the dirichlet calibration layer
      self.dirichlet_linear.weight.data.copy_(
          torch.eye(self.dirichlet_linear.weight.shape[0])
      )
      self.dirichlet_linear.bias.data.copy_(
          torch.zeros(*self.dirichlet_linear.bias.shape)
      )
      if use_cuda:
          self.dirichlet_linear.cuda()
  
      optimizer = Adam(self.dirichlet_linear.parameters(), lr=self.lr)
  
      loss_history, weights = self.wrapper.train_and_test_on_datasets(
          train_set,
          test_set,
          optimizer,
          batch_size,
          epoch,
          use_cuda,
          regularizer=self.l2_reg,
          return_best_weights=True,
          patience=None,
          **kwargs,
      )
      self.model.load_state_dict(weights)
  
      if double_fit:
          lr = self.lr / 10
          optimizer = Adam(self.dirichlet_linear.parameters(), lr=lr)
          (
              loss_history,
              weights,
          ) = self.wrapper.train_and_test_on_datasets(
              train_set,
              test_set,
              optimizer,
              batch_size,
              epoch,
              use_cuda,
              regularizer=self.l2_reg,
              return_best_weights=True,
              patience=None,
              **kwargs,
          )
          self.model.load_state_dict(weights)
  
      return loss_history, self.model.state_dict()
  
file_path: baal/baal/calibration/calibration.py
imports: from copy import deepcopy
  import structlog
  import torch
  from baal.active.dataset.base import Dataset
  from torch import nn
  from torch.optim import Adam
  from baal import ModelWrapper
  from baal.utils.metrics import ECE, ECE_PerCLs
label_desc: Use image node_id 3 for calling the DirichletCalibrator obj's underlying member method code with example usage: obj.calibrate(train_set, test_set, batch_size, epoch, use_cuda, double_fit) and returns: loss_history,  self