code: def initialize_optimizer(
      model,
      n_batches,
      n_epochs,
      device,
      learning_rate,
      optimizer_opts=None,
      schedule_opts=None,
      distributed=False,
      grad_acc_steps=1,
      local_rank=-1,
      use_amp=None,
  ):
      """
      Initializes an optimizer, a learning rate scheduler and converts the model if needed (e.g for mixed precision).
      Per default, we use transformers' AdamW and a linear warmup schedule with warmup ratio 0.1.
      You can easily switch optimizer and schedule via `optimizer_opts` and `schedule_opts`.
  
      :param model: model to optimize (e.g. trimming weights to fp16 / mixed precision)
      :type model: AdaptiveModel
      :param n_batches: number of batches for training
      :type n_batches: int
      :param n_epochs: number of epochs for training
      :param device:
      :param learning_rate: Learning rate
      :type learning_rate: float
      :param optimizer_opts: Dict to customize the optimizer. Choose any optimizer available from torch.optim, apex.optimizers or
                             transformers.optimization by supplying the class name and the parameters for the constructor.
                             Examples:
                             1) AdamW from Transformers (Default):
                             {"name": "TransformersAdamW", "correct_bias": False, "weight_decay": 0.01}
                             2) SGD from pytorch:
                             {"name": "SGD", "momentum": 0.0}
                             3) FusedLAMB from apex:
                             {"name": "FusedLAMB", "bias_correction": True}
      :param schedule_opts: Dict to customize the learning rate schedule.
                            Choose any Schedule from Pytorch or Huggingface's Transformers by supplying the class name
                            and the parameters needed by the constructor.
                            If the dict does not contain ``num_training_steps`` it will be set by
                            calculating it from ``n_batches``, ``grad_acc_steps`` and ``n_epochs``.
                            Examples:
                            1) Linear Warmup (Default):
                            {"name": "LinearWarmup",
                            "num_warmup_steps": 0.1 * num_training_steps,
                            "num_training_steps": num_training_steps}
                            2) CosineWarmup:
                            {"name": "CosineWarmup",
                            "num_warmup_steps": 0.1 * num_training_steps,
                            "num_training_steps": num_training_steps}
                            3) CyclicLR from pytorch:
                            {"name": "CyclicLR", "base_lr": 1e-5, "max_lr":1e-4, "step_size_up": 100}
      :param distributed: Whether training on distributed machines
      :param grad_acc_steps: Number of steps to accumulate gradients for. Helpful to mimic large batch_sizes on small machines.
      :param local_rank: rank of the machine in a distributed setting
      :param use_amp: Optimization level of nvidia's automatic mixed precision (AMP). The higher the level, the faster the model.
                      Options:
                      "O0" (Normal FP32 training)
                      "O1" (Mixed Precision => Recommended)
                      "O2" (Almost FP16)
                      "O3" (Pure FP16).
                      See details on: https://nvidia.github.io/apex/amp.html
      :return: model, optimizer, scheduler
      """
  
      if use_amp and not AMP_AVAILABLE:
          raise ImportError(
              f"Got use_amp = {use_amp}, but cannot find apex. "
              "Please install Apex if you want to make use of automatic mixed precision. "
              "https://github.com/NVIDIA/apex"
          )
  
      if (schedule_opts is not None) and (
          not isinstance(schedule_opts, dict)
      ):
          raise TypeError(
              "Parameter schedule_opts must be None or "
              "an instance of dict but was {}!".format(
                  type(schedule_opts)
              )
          )
  
      num_train_optimization_steps = (
          int(n_batches / grad_acc_steps) * n_epochs
      )
  
      # Use some defaults to simplify life of inexperienced users
      if optimizer_opts is None:
          optimizer_opts = {
              "name": "TransformersAdamW",
              "correct_bias": False,
              "weight_decay": 0.01,
          }
      optimizer_opts["lr"] = learning_rate
  
      if schedule_opts is None:
          # Default schedule: Linear Warmup with 10% warmup
          schedule_opts = {
              "name": "LinearWarmup",
              "num_warmup_steps": 0.1 * num_train_optimization_steps,
              "num_training_steps": num_train_optimization_steps,
          }
  
          # schedule_opts = {"name": "OneCycleLR", "max_lr":learning_rate, "pct_start": 0.1,
          #                  "total_steps": num_train_optimization_steps }
      elif "num_training_steps" not in schedule_opts:
          schedule_opts[
              "num_training_steps"
          ] = num_train_optimization_steps
  
      # Log params
      MlLogger.log_params(
          {
              "use_amp": use_amp,
              "num_train_optimization_steps": schedule_opts[
                  "num_training_steps"
              ],
          }
      )
  
      # Get optimizer from pytorch, transformers or apex
      optimizer = _get_optim(model, optimizer_opts)
  
      # Adjust for parallel training + amp
      model, optimizer = optimize_model(
          model, device, local_rank, optimizer, distributed, use_amp
      )
  
      # Get learning rate schedule - moved below to supress warning
      scheduler = get_scheduler(optimizer, schedule_opts)
  
      return model, optimizer, scheduler
  
file_path: FARM/farm/modeling/optimization.py
imports: from importlib import import_module
  import logging
  import sys
  import inspect
  import torch
  from torch.nn.parallel import DistributedDataParallel
  from torch.nn import DataParallel
  from transformers.optimization import AdamW
  from farm.utils import MLFlowLogger
label_desc: Use image node_id 1 for calling a global function with example usage: initialize_optimizer(model, n_batches, n_epochs, device, learning_rate, optimizer_opts, schedule_opts, distributed, grad_acc_steps, local_rank, use_amp) and returns: model,  optimizer,  scheduler