docstr: Train the kernel to maximise an estimate of test power using minibatch gradient descent.
docstr_tok: ["Train","the","kernel","to","maximise","an","estimate","of","test","power","using","minibatch","gradient","descent","."]
code: def trainer(
      j_hat: JHat,
      datasets: Tuple[tf.keras.utils.Sequence, tf.keras.utils.Sequence],
      optimizer: OptimizerTF = tf.keras.optimizers.Adam,
      learning_rate: float = 1e-3,
      preprocess_fn: Callable = None,
      epochs: int = 20,
      reg_loss_fn: Callable = (lambda kernel: 0),
      verbose: int = 1,
  ) -> None:
      """
      Train the kernel to maximise an estimate of test power using minibatch gradient descent.
      """
      ds_ref, ds_cur = datasets
      optimizer = (
          optimizer(learning_rate=learning_rate)
          if isinstance(optimizer, type)
          else optimizer
      )
      n_minibatch = min(len(ds_ref), len(ds_cur))
      # iterate over epochs
      loss_ma = 0.0
      for epoch in range(epochs):
          if verbose:
              pbar = tf.keras.utils.Progbar(n_minibatch, 1)
          for step, (x_ref, x_cur) in enumerate(zip(ds_ref, ds_cur)):
              if isinstance(preprocess_fn, Callable):  # type: ignore
                  x_ref, x_cur = preprocess_fn(x_ref), preprocess_fn(
                      x_cur
                  )
              with tf.GradientTape() as tape:
                  estimate = j_hat(x_ref, x_cur)
                  loss = -estimate + reg_loss_fn(j_hat.kernel)  # ascent
              grads = tape.gradient(loss, j_hat.trainable_weights)
              optimizer.apply_gradients(
                  zip(grads, j_hat.trainable_weights)
              )
              if verbose == 1:
                  loss_ma = loss_ma + (loss - loss_ma) / (step + 1)
                  pbar_values = [("loss", loss_ma)]
                  pbar.add(1, values=pbar_values)
  
file_path: alibi-detect/alibi_detect/cd/tensorflow/learned_kernel.py
imports: from functools import partial
  import numpy
  import tensorflow
  from typing import Callable, Dict, Optional, Tuple, Union
  from alibi_detect.cd.base import BaseLearnedKernelDrift
  from alibi_detect.utils.tensorflow.data import TFDataset
  from alibi_detect.utils.tensorflow.misc import clone_model
  from alibi_detect.utils.tensorflow.distance import mmd2_from_kernel_matrix, batch_compute_kernel_matrix
  from alibi_detect.utils.warnings import deprecated_alias
  from alibi_detect.utils.frameworks import Framework
  from alibi_detect.utils._types import OptimizerTF
label_desc: Use image node_id 3 for calling the LearnedKernelDriftTF obj's underlying member method code with example usage: obj.trainer(j_hat, datasets, optimizer, learning_rate, preprocess_fn, epochs, reg_loss_fn, verbose) without return types