code: def train(args):
      print("start training...")
  
      """@nni.variable(nni.choice('UNetResNetV4', 'UNetResNetV5', 'UNetResNetV6'), name=model_name)"""
      model_name = args.model_name
  
      model = eval(model_name)(args.layers, num_filters=args.nf)
      model_subdir = args.pad_mode
      if args.meta_version == 2:
          model_subdir = args.pad_mode + "_meta2"
      if args.exp_name is None:
          model_file = os.path.join(
              MODEL_DIR,
              model.name,
              model_subdir,
              "best_{}.pth".format(args.ifold),
          )
      else:
          model_file = os.path.join(
              MODEL_DIR,
              args.exp_name,
              model.name,
              model_subdir,
              "best_{}.pth".format(args.ifold),
          )
  
      parent_dir = os.path.dirname(model_file)
      if not os.path.exists(parent_dir):
          os.makedirs(parent_dir)
  
      if args.init_ckp is not None:
          CKP = args.init_ckp
      else:
          CKP = model_file
      if os.path.exists(CKP):
          print("loading {}...".format(CKP))
          model.load_state_dict(torch.load(CKP))
      model = model.cuda()
  
      if args.optim == "Adam":
          optimizer = optim.Adam(
              model.parameters(), lr=args.lr, weight_decay=0.0001
          )
      else:
          optimizer = optim.SGD(
              model.parameters(),
              lr=args.lr,
              momentum=0.9,
              weight_decay=0.0001,
          )
  
      train_loader, val_loader = get_train_loaders(
          args.ifold,
          batch_size=args.batch_size,
          dev_mode=args.dev_mode,
          pad_mode=args.pad_mode,
          meta_version=args.meta_version,
          pseudo_label=args.pseudo,
          depths=args.depths,
      )
  
      if args.lrs == "plateau":
          lr_scheduler = ReduceLROnPlateau(
              optimizer,
              mode="max",
              factor=args.factor,
              patience=args.patience,
              min_lr=args.min_lr,
          )
      else:
          lr_scheduler = CosineAnnealingLR(
              optimizer, args.t_max, eta_min=args.min_lr
          )
  
      print(
          "epoch |   lr    |   %       |  loss  |  avg   | f loss | lovaz  |  iou   | iout   |  best  | time | save |  salt  |"
      )
  
      best_iout, _iou, _f, _l, _salt, best_mix_score = validate(
          args, model, val_loader, args.start_epoch
      )
      print(
          "val   |         |           |        |        | {:.4f} | {:.4f} | {:.4f} | {:.4f} | {:.4f} |      |      | {:.4f} |".format(
              _f, _l, _iou, best_iout, best_iout, _salt
          )
      )
      if args.val:
          return
  
      model.train()
  
      if args.lrs == "plateau":
          lr_scheduler.step(best_iout)
      else:
          lr_scheduler.step()
  
      for epoch in range(args.start_epoch, args.epochs):
          train_loss = 0
  
          current_lr = get_lrs(optimizer)
          bg = time.time()
          for batch_idx, data in enumerate(train_loader):
              img, target, salt_target = data
              if args.depths:
                  add_depth_channel(img, args.pad_mode)
              img, target, salt_target = (
                  img.cuda(),
                  target.cuda(),
                  salt_target.cuda(),
              )
              optimizer.zero_grad()
              output, salt_out = model(img)
  
              loss, *_ = weighted_loss(
                  args,
                  (output, salt_out),
                  (target, salt_target),
                  epoch=epoch,
              )
              loss.backward()
  
              if args.optim == "Adam" and args.adamw:
                  wd = 0.0001
                  for group in optimizer.param_groups:
                      for param in group["params"]:
                          param.data = param.data.add(
                              -wd * group["lr"], param.data
                          )
  
              optimizer.step()
  
              train_loss += loss.item()
              print(
                  "\r {:4d} | {:.5f} | {:4d}/{} | {:.4f} | {:.4f} |".format(
                      epoch,
                      float(current_lr[0]),
                      args.batch_size * (batch_idx + 1),
                      train_loader.num,
                      loss.item(),
                      train_loss / (batch_idx + 1),
                  ),
                  end="",
              )
  
          (
              iout,
              iou,
              focal_loss,
              lovaz_loss,
              salt_loss,
              mix_score,
          ) = validate(args, model, val_loader, epoch=epoch)
          """@nni.report_intermediate_result(iout)"""
  
          _save_ckp = ""
          if iout > best_iout:
              best_iout = iout
              torch.save(model.state_dict(), model_file)
              _save_ckp = "*"
          if args.store_loss_model and mix_score > best_mix_score:
              best_mix_score = mix_score
              torch.save(model.state_dict(), model_file + "_loss")
              _save_ckp += "."
          print(
              " {:.4f} | {:.4f} | {:.4f} | {:.4f} | {:.4f} | {:.2f} | {:4s} | {:.4f} |".format(
                  focal_loss,
                  lovaz_loss,
                  iou,
                  iout,
                  best_iout,
                  (time.time() - bg) / 60,
                  _save_ckp,
                  salt_loss,
              )
          )
  
          model.train()
  
          if args.lrs == "plateau":
              lr_scheduler.step(best_iout)
          else:
              lr_scheduler.step()
  
      del model, train_loader, val_loader, optimizer, lr_scheduler
      """@nni.report_final_result(best_iout)"""
  
imports: import os
  import argparse
  import time
  import torch
  import torch.nn
  import torch.optim
  import torch.nn.functional
  from torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau
  from loader import get_train_loaders, add_depth_channel
  from models import UNetResNetV4, UNetResNetV5, UNetResNetV6
  from lovasz_losses import lovasz_hinge
  from focal_loss import FocalLoss2d
  from postprocessing import binarize, crop_image, resize_image
  from metrics import intersection_over_union, intersection_over_union_thresholds
  import settings
label_desc: Use image node_id 2 for calling a global function with example usage: train(args) without return types