code: def train(self, buffer_dict):
      torch.cuda.set_device(self.gpu)
  
      # DATASET -----------------------------------------------------------------------------------------------------#
      trn_loader = self.trn_loader
  
      # MODEL -------------------------------------------------------------------------------------------------------#
      if self.config.train.backbone == "vgg":
          craft = CRAFT(pretrained=True, amp=self.config.train.amp)
      else:
          raise Exception("Undefined architecture")
  
      if self.config.train.ckpt_path is not None:
          craft.load_state_dict(copyStateDict(self.net_param["craft"]))
      craft = nn.SyncBatchNorm.convert_sync_batchnorm(craft)
      craft = craft.cuda()
      craft = torch.nn.parallel.DistributedDataParallel(
          craft, device_ids=[self.gpu]
      )
  
      torch.backends.cudnn.benchmark = True
  
      # OPTIMIZER----------------------------------------------------------------------------------------------------#
  
      optimizer = optim.Adam(
          craft.parameters(),
          lr=self.config.train.lr,
          weight_decay=self.config.train.weight_decay,
      )
  
      if (
          self.config.train.ckpt_path is not None
          and self.config.train.st_iter != 0
      ):
          optimizer.load_state_dict(
              copyStateDict(self.net_param["optimizer"])
          )
          self.config.train.st_iter = self.net_param["optimizer"][
              "state"
          ][0]["step"]
          self.config.train.lr = self.net_param["optimizer"][
              "param_groups"
          ][0]["lr"]
  
      # LOSS --------------------------------------------------------------------------------------------------------#
      # mixed precision
      if self.config.train.amp:
          scaler = torch.cuda.amp.GradScaler()
  
          # load model
          if (
              self.config.train.ckpt_path is not None
              and self.config.train.st_iter != 0
          ):
              scaler.load_state_dict(
                  copyStateDict(self.net_param["scaler"])
              )
      else:
          scaler = None
  
      criterion = self.get_loss()
  
      # TRAIN -------------------------------------------------------------------------------------------------------#
      train_step = self.config.train.st_iter
      whole_training_step = self.config.train.end_iter
      update_lr_rate_step = 0
      training_lr = self.config.train.lr
      loss_value = 0
      batch_time = 0
      epoch = 0
      start_time = time.time()
  
      while train_step < whole_training_step:
          self.trn_sampler.set_epoch(train_step)
          for (
              index,
              (
                  image,
                  region_image,
                  affinity_image,
                  confidence_mask,
              ),
          ) in enumerate(trn_loader):
              craft.train()
              if (
                  train_step > 0
                  and train_step % self.config.train.lr_decay == 0
              ):
                  update_lr_rate_step += 1
                  training_lr = self.adjust_learning_rate(
                      optimizer,
                      self.config.train.gamma,
                      update_lr_rate_step,
                      self.config.train.lr,
                  )
  
              images = image.cuda(non_blocking=True)
              region_image_label = region_image.cuda(non_blocking=True)
              affinity_image_label = affinity_image.cuda(
                  non_blocking=True
              )
              confidence_mask_label = confidence_mask.cuda(
                  non_blocking=True
              )
  
              if self.config.train.amp:
                  with torch.cuda.amp.autocast():
                      output, _ = craft(images)
                      out1 = output[:, :, :, 0]
                      out2 = output[:, :, :, 1]
  
                      loss = criterion(
                          region_image_label,
                          affinity_image_label,
                          out1,
                          out2,
                          confidence_mask_label,
                          self.config.train.neg_rto,
                          self.config.train.n_min_neg,
                      )
  
                  optimizer.zero_grad()
                  scaler.scale(loss).backward()
                  scaler.step(optimizer)
                  scaler.update()
  
              else:
                  output, _ = craft(images)
                  out1 = output[:, :, :, 0]
                  out2 = output[:, :, :, 1]
                  loss = criterion(
                      region_image_label,
                      affinity_image_label,
                      out1,
                      out2,
                      confidence_mask_label,
                      self.config.train.neg_rto,
                  )
  
                  optimizer.zero_grad()
                  loss.backward()
                  optimizer.step()
  
              end_time = time.time()
              loss_value += loss.item()
              batch_time += end_time - start_time
  
              if (
                  train_step > 0
                  and train_step % 5 == 0
                  and self.gpu == 0
              ):
                  mean_loss = loss_value / 5
                  loss_value = 0
                  avg_batch_time = batch_time / 5
                  batch_time = 0
  
                  print(
                      "{}, training_step: {}|{}, learning rate: {:.8f}, "
                      "training_loss: {:.5f}, avg_batch_time: {:.5f}".format(
                          time.strftime(
                              "%Y-%m-%d:%H:%M:%S",
                              time.localtime(time.time()),
                          ),
                          train_step,
                          whole_training_step,
                          training_lr,
                          mean_loss,
                          avg_batch_time,
                      )
                  )
                  if self.gpu == 0 and self.config.wandb_opt:
                      wandb.log(
                          {
                              "train_step": train_step,
                              "mean_loss": mean_loss,
                          }
                      )
  
              if (
                  train_step % self.config.train.eval_interval == 0
                  and train_step != 0
              ):
                  # initialize all buffer value with zero
                  if self.gpu == 0:
                      for buffer in buffer_dict.values():
                          for i in range(len(buffer)):
                              buffer[i] = None
  
                  print("Saving state, index:", train_step)
                  save_param_dic = {
                      "iter": train_step,
                      "craft": craft.state_dict(),
                      "optimizer": optimizer.state_dict(),
                  }
                  save_param_path = (
                      self.config.results_dir
                      + "/CRAFT_clr_"
                      + repr(train_step)
                      + ".pth"
                  )
  
                  if self.config.train.amp:
                      save_param_dic["scaler"] = scaler.state_dict()
                      save_param_path = (
                          self.config.results_dir
                          + "/CRAFT_clr_amp_"
                          + repr(train_step)
                          + ".pth"
                      )
  
                  if self.gpu == 0:
                      torch.save(save_param_dic, save_param_path)
  
                  # validation
                  self.iou_eval(
                      "icdar2013",
                      train_step,
                      save_param_path,
                      buffer_dict["icdar2013"],
                      craft,
                  )
  
              train_step += 1
              if train_step >= whole_training_step:
                  break
          epoch += 1
  
      # save last model
      if self.gpu == 0:
          save_param_dic = {
              "iter": train_step,
              "craft": craft.state_dict(),
              "optimizer": optimizer.state_dict(),
          }
          save_param_path = (
              self.config.results_dir
              + "/CRAFT_clr_"
              + repr(train_step)
              + ".pth"
          )
  
          if self.config.train.amp:
              save_param_dic["scaler"] = scaler.state_dict()
              save_param_path = (
                  self.config.results_dir
                  + "/CRAFT_clr_amp_"
                  + repr(train_step)
                  + ".pth"
              )
          torch.save(save_param_dic, save_param_path)
  
imports: import argparse
  import os
  import shutil
  import time
  import yaml
  import multiprocessing
  import numpy
  import torch
  import torch.nn
  import torch.optim
  import wandb
  from config.load_config import load_yaml, DotDict
  from data.dataset import SynthTextDataSet
  from loss.mseloss import Maploss_v2, Maploss_v3
  from model.craft import CRAFT
  from metrics.eval_det_iou import DetectionIoUEvaluator
  from eval import main_eval
  from utils.util import copyStateDict, save_parser
label_desc: Use image node_id 7 for calling the Trainer obj's underlying member method code with example usage: obj.train(buffer_dict) without return types