code: def train_model(
      args,
      model,
      loaders,
      *,
      checkpoint=None,
      dp_device_ids=None,
      store=None,
      update_params=None,
      disable_no_grad=False,
  ):
      """
      Main function for training a model.
  
      Args:
          args (object) : A python object for arguments, implementing
              ``getattr()`` and ``setattr()`` and having the following
              attributes. See :attr:`robustness.defaults.TRAINING_ARGS` for a
              list of arguments, and you can use
              :meth:`robustness.defaults.check_and_fill_args` to make sure that
              all required arguments are filled and to fill missing args with
              reasonable defaults:
  
              adv_train (int or bool, *required*)
                  if 1/True, adversarially train, otherwise if 0/False do
                  standard training
              epochs (int, *required*)
                  number of epochs to train for
              lr (float, *required*)
                  learning rate for SGD optimizer
              weight_decay (float, *required*)
                  weight decay for SGD optimizer
              momentum (float, *required*)
                  momentum parameter for SGD optimizer
              step_lr (int)
                  if given, drop learning rate by 10x every `step_lr` steps
              custom_lr_multplier (str)
                  If given, use a custom LR schedule, formed by multiplying the
                      original ``lr`` (format: [(epoch, LR_MULTIPLIER),...])
              lr_interpolation (str)
                  How to drop the learning rate, either ``step`` or ``linear``,
                      ignored unless ``custom_lr_multiplier`` is provided.
              adv_eval (int or bool)
                  If True/1, then also do adversarial evaluation, otherwise skip
                  (ignored if adv_train is True)
              log_iters (int, *required*)
                  How frequently (in epochs) to save training logs
              save_ckpt_iters (int, *required*)
                  How frequently (in epochs) to save checkpoints (if -1, then only
                  save latest and best ckpts)
              attack_lr (float or str, *required if adv_train or adv_eval*)
                  float (or float-parseable string) for the adv attack step size
              constraint (str, *required if adv_train or adv_eval*)
                  the type of adversary constraint
                  (:attr:`robustness.attacker.STEPS`)
              eps (float or str, *required if adv_train or adv_eval*)
                  float (or float-parseable string) for the adv attack budget
              attack_steps (int, *required if adv_train or adv_eval*)
                  number of steps to take in adv attack
              custom_eps_multiplier (str, *required if adv_train or adv_eval*)
                  If given, then set epsilon according to a schedule by
                  multiplying the given eps value by a factor at each epoch. Given
                  in the same format as ``custom_lr_multiplier``, ``[(epoch,
                  MULTIPLIER)..]``
              use_best (int or bool, *required if adv_train or adv_eval*) :
                  If True/1, use the best (in terms of loss) PGD step as the
                  attack, if False/0 use the last step
              random_restarts (int, *required if adv_train or adv_eval*)
                  Number of random restarts to use for adversarial evaluation
              custom_train_loss (function, optional)
                  If given, a custom loss instead of the default CrossEntropyLoss.
                  Takes in `(logits, targets)` and returns a scalar.
              custom_adv_loss (function, *required if custom_train_loss*)
                  If given, a custom loss function for the adversary. The custom
                  loss function takes in `model, input, target` and should return
                  a vector representing the loss for each element of the batch, as
                  well as the classifier output.
              custom_accuracy (function)
                  If given, should be a function that takes in model outputs
                  and model targets and outputs a top1 and top5 accuracy, will
                  displayed instead of conventional accuracies
              regularizer (function, optional)
                  If given, this function of `model, input, target` returns a
                  (scalar) that is added on to the training loss without being
                  subject to adversarial attack
              iteration_hook (function, optional)
                  If given, this function is called every training iteration by
                  the training loop (useful for custom logging). The function is
                  given arguments `model, iteration #, loop_type [train/eval],
                  current_batch_ims, current_batch_labels`.
              epoch hook (function, optional)
                  Similar to iteration_hook but called every epoch instead, and
                  given arguments `model, log_info` where `log_info` is a
                  dictionary with keys `epoch, nat_prec1, adv_prec1, nat_loss,
                  adv_loss, train_prec1, train_loss`.
  
          model (AttackerModel) : the model to train.
          loaders (tuple[iterable]) : `tuple` of data loaders of the form
              `(train_loader, val_loader)`
          checkpoint (dict) : a loaded checkpoint previously saved by this library
              (if resuming from checkpoint)
          dp_device_ids (list|None) : if not ``None``, a list of device ids to
              use for DataParallel.
          store (cox.Store) : a cox store for logging training progress
          update_params (list) : list of parameters to use for training, if None
              then all parameters in the model are used (useful for transfer
              learning)
          disable_no_grad (bool) : if True, then even model evaluation will be
              run with autograd enabled (otherwise it will be wrapped in a ch.no_grad())
      """
      # Logging setup
      writer = store.tensorboard if store else None
      prec1_key = f"{'adv' if args.adv_train else 'nat'}_prec1"
      if store is not None:
          store.add_table(consts.LOGS_TABLE, consts.LOGS_SCHEMA)
  
      # Reformat and read arguments
      check_required_args(args)  # Argument sanity check
      for p in ["eps", "attack_lr", "custom_eps_multiplier"]:
          setattr(
              args,
              p,
              eval(str(getattr(args, p)))
              if has_attr(args, p)
              else None,
          )
      if args.custom_eps_multiplier is not None:
          eps_periods = args.custom_eps_multiplier
          args.custom_eps_multiplier = lambda t: np.interp(
              [t], *zip(*eps_periods)
          )[0]
  
      # Initial setup
      train_loader, val_loader = loaders
      opt, schedule = make_optimizer_and_schedule(
          args, model, checkpoint, update_params
      )
  
      # Put the model into parallel mode
      assert not hasattr(
          model, "module"
      ), "model is already in DataParallel."
      model = ch.nn.DataParallel(model, device_ids=dp_device_ids).cuda()
  
      best_prec1, start_epoch = (0, 0)
      if checkpoint:
          start_epoch = checkpoint["epoch"]
          best_prec1 = (
              checkpoint[prec1_key]
              if prec1_key in checkpoint
              else _model_loop(
                  args,
                  "val",
                  val_loader,
                  model,
                  None,
                  start_epoch - 1,
                  args.adv_train,
                  writer=None,
              )[0]
          )
  
      # Timestamp for training start time
      start_time = time.time()
  
      for epoch in range(start_epoch, args.epochs):
          # train for one epoch
          train_prec1, train_loss = _model_loop(
              args,
              "train",
              train_loader,
              model,
              opt,
              epoch,
              args.adv_train,
              writer,
          )
          last_epoch = epoch == (args.epochs - 1)
  
          # evaluate on validation set
          sd_info = {
              "model": model.state_dict(),
              "optimizer": opt.state_dict(),
              "schedule": (schedule and schedule.state_dict()),
              "epoch": epoch + 1,
              "amp": amp.state_dict() if args.mixed_precision else None,
          }
  
          def save_checkpoint(filename):
              ckpt_save_path = os.path.join(
                  args.out_dir if not store else store.path, filename
              )
              ch.save(sd_info, ckpt_save_path, pickle_module=dill)
  
          save_its = args.save_ckpt_iters
          should_save_ckpt = (epoch % save_its == 0) and (save_its > 0)
          should_log = epoch % args.log_iters == 0
  
          if should_log or last_epoch or should_save_ckpt:
              # log + get best
              ctx = (
                  ch.enable_grad() if disable_no_grad else ch.no_grad()
              )
              with ctx:
                  prec1, nat_loss = _model_loop(
                      args,
                      "val",
                      val_loader,
                      model,
                      None,
                      epoch,
                      False,
                      writer,
                  )
  
              # loader, model, epoch, input_adv_exs
              should_adv_eval = args.adv_eval or args.adv_train
              adv_val = should_adv_eval and _model_loop(
                  args,
                  "val",
                  val_loader,
                  model,
                  None,
                  epoch,
                  True,
                  writer,
              )
              adv_prec1, adv_loss = adv_val or (-1.0, -1.0)
  
              # remember best prec@1 and save checkpoint
              our_prec1 = adv_prec1 if args.adv_train else prec1
              is_best = our_prec1 > best_prec1
              best_prec1 = max(our_prec1, best_prec1)
              sd_info[prec1_key] = our_prec1
  
              # log every checkpoint
              log_info = {
                  "epoch": epoch + 1,
                  "nat_prec1": prec1,
                  "adv_prec1": adv_prec1,
                  "nat_loss": nat_loss,
                  "adv_loss": adv_loss,
                  "train_prec1": train_prec1,
                  "train_loss": train_loss,
                  "time": time.time() - start_time,
              }
  
              # Log info into the logs table
              if store:
                  store[consts.LOGS_TABLE].append_row(log_info)
              # If we are at a saving epoch (or the last epoch), save a checkpoint
              if should_save_ckpt or last_epoch:
                  save_checkpoint(ckpt_at_epoch(epoch))
  
              # Update the latest and best checkpoints (overrides old one)
              save_checkpoint(consts.CKPT_NAME_LATEST)
              if is_best:
                  save_checkpoint(consts.CKPT_NAME_BEST)
  
          if schedule:
              schedule.step()
          if has_attr(args, "epoch_hook"):
              args.epoch_hook(model, log_info)
  
      return model
  
file_path: robustness/robustness/train.py
imports: import torch
  import numpy
  import torch.nn
  from torch.optim import SGD, lr_scheduler
  from torchvision.utils import make_grid
  from cox.utils import Parameters
  from .tools import helpers
  from .tools.helpers import AverageMeter, ckpt_at_epoch, has_attr
  from .tools import constants
  import dill
  import os
  import time
  import warnings
label_desc: Use image node_id 4 for calling a global function with example usage: train_model(args, model, loaders) and returns: model