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,
      dataloaders: Tuple[DataLoader, DataLoader],
      device: torch.device,
      optimizer: Callable = torch.optim.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.
      """
      optimizer = optimizer(j_hat.parameters(), lr=learning_rate)
      j_hat.train()
      loss_ma = 0.0
      for epoch in range(epochs):
          dl_ref, dl_cur = dataloaders
          dl = (
              tqdm(
                  enumerate(zip(dl_ref, dl_cur)),
                  total=min(len(dl_ref), len(dl_cur)),
              )
              if verbose == 1
              else enumerate(zip(dl_ref, dl_cur))
          )
          for step, (x_ref, x_cur) in dl:
              if isinstance(preprocess_fn, Callable):  # type: ignore
                  x_ref, x_cur = preprocess_fn(x_ref), preprocess_fn(
                      x_cur
                  )
              x_ref, x_cur = x_ref.to(device), x_cur.to(device)
              optimizer.zero_grad()  # type: ignore
              estimate = j_hat(x_ref, x_cur)
              loss = -estimate + reg_loss_fn(j_hat.kernel)  # ascent
              loss.backward()
              optimizer.step()  # type: ignore
              if verbose == 1:
                  loss_ma = loss_ma + (loss.item() - loss_ma) / (
                      step + 1
                  )
                  dl.set_description(f"Epoch {epoch + 1}/{epochs}")
                  dl.set_postfix(dict(loss=loss_ma))
  
file_path: alibi-detect/alibi_detect/cd/pytorch/learned_kernel.py
imports: from copy import deepcopy
  from functools import partial
  from tqdm import tqdm
  import numpy
  import torch
  import torch.nn
  from torch.utils.data import DataLoader
  from typing import Callable, Dict, Optional, Union, Tuple
  from alibi_detect.cd.base import BaseLearnedKernelDrift
  from alibi_detect.utils.pytorch import get_device
  from alibi_detect.utils.pytorch.distance import mmd2_from_kernel_matrix, batch_compute_kernel_matrix
  from alibi_detect.utils.pytorch.data import TorchDataset
  from alibi_detect.utils.warnings import deprecated_alias
  from alibi_detect.utils.frameworks import Framework
  from alibi_detect.utils._types import TorchDeviceType
label_desc: Use image node_id 3 for calling the LearnedKernelDriftTorch obj's underlying member method code with example usage: obj.trainer(j_hat, dataloaders, device, optimizer, learning_rate, preprocess_fn, epochs, reg_loss_fn, verbose) without return types