code: def train(self, buffer_dict):
      torch.cuda.set_device(self.gpu)
      total_gpu_num = torch.cuda.device_count()
  
      # MODEL -------------------------------------------------------------------------------------------------------#
      # SUPERVISION model
      if self.config.mode == "weak_supervision":
          if self.config.train.backbone == "vgg":
              supervision_model = CRAFT(
                  pretrained=False, amp=self.config.train.amp
              )
          else:
              raise Exception("Undefined architecture")
  
          # NOTE: only work on half GPU assign train / half GPU assign supervision setting
          supervision_device = total_gpu_num // 2 + self.gpu
          if self.config.train.ckpt_path is not None:
              supervision_param = self.get_load_param(
                  supervision_device
              )
              supervision_model.load_state_dict(
                  copyStateDict(supervision_param["craft"])
              )
              supervision_model = supervision_model.to(
                  f"cuda:{supervision_device}"
              )
          print(
              f"Supervision model loading on : gpu {supervision_device}"
          )
      else:
          supervision_model, supervision_device = None, None
  
      # TRAIN model
      if self.config.train.backbone == "vgg":
          craft = CRAFT(pretrained=False, 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
  
      # DATASET -----------------------------------------------------------------------------------------------------#
  
      if self.config.train.use_synthtext:
          trn_syn_loader = self.get_synth_loader()
          batch_syn = iter(trn_syn_loader)
  
      if self.config.train.real_dataset == "custom":
          trn_real_dataset = self.get_custom_dataset()
      else:
          raise Exception("Undefined dataset")
  
      if self.config.mode == "weak_supervision":
          trn_real_dataset.update_model(supervision_model)
          trn_real_dataset.update_device(supervision_device)
  
      trn_real_sampler = (
          torch.utils.data.distributed.DistributedSampler(
              trn_real_dataset
          )
      )
      trn_real_loader = torch.utils.data.DataLoader(
          trn_real_dataset,
          batch_size=self.config.train.batch_size,
          shuffle=False,
          num_workers=self.config.train.num_workers,
          sampler=trn_real_sampler,
          drop_last=False,
          pin_memory=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()
  
          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
      start_time = time.time()
  
      print(
          "================================ Train start ================================"
      )
      while train_step < whole_training_step:
          trn_real_sampler.set_epoch(train_step)
          for (
              index,
              (
                  images,
                  region_scores,
                  affinity_scores,
                  confidence_masks,
              ),
          ) in enumerate(trn_real_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 = images.cuda(non_blocking=True)
              region_scores = region_scores.cuda(non_blocking=True)
              affinity_scores = affinity_scores.cuda(non_blocking=True)
              confidence_masks = confidence_masks.cuda(
                  non_blocking=True
              )
  
              if self.config.train.use_synthtext:
                  # Synth image load
                  (
                      syn_image,
                      syn_region_label,
                      syn_affi_label,
                      syn_confidence_mask,
                  ) = next(batch_syn)
                  syn_image = syn_image.cuda(non_blocking=True)
                  syn_region_label = syn_region_label.cuda(
                      non_blocking=True
                  )
                  syn_affi_label = syn_affi_label.cuda(
                      non_blocking=True
                  )
                  syn_confidence_mask = syn_confidence_mask.cuda(
                      non_blocking=True
                  )
  
                  # concat syn & custom image
                  images = torch.cat((syn_image, images), 0)
                  region_image_label = torch.cat(
                      (syn_region_label, region_scores), 0
                  )
                  affinity_image_label = torch.cat(
                      (syn_affi_label, affinity_scores), 0
                  )
                  confidence_mask_label = torch.cat(
                      (syn_confidence_mask, confidence_masks), 0
                  )
              else:
                  region_image_label = region_scores
                  affinity_image_label = affinity_scores
                  confidence_mask_label = confidence_masks
  
              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
              ):
                  craft.eval()
                  # 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"
                          )
  
                      torch.save(save_param_dic, save_param_path)
  
                  # validation
                  self.iou_eval(
                      "custom_data",
                      train_step,
                      buffer_dict["custom_data"],
                      craft,
                  )
  
              train_step += 1
              if train_step >= whole_training_step:
                  break
  
          if self.config.mode == "weak_supervision":
              state_dict = craft.module.state_dict()
              supervision_model.load_state_dict(state_dict)
              trn_real_dataset.update_model(supervision_model)
  
      # 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 multiprocessing
  import yaml
  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, CustomDataset
  from loss.mseloss import Maploss_v2, Maploss_v3
  from model.craft import CRAFT
  from eval import main_eval
  from metrics.eval_det_iou import DetectionIoUEvaluator
  from utils.util import copyStateDict, save_parser
label_desc: Use image node_id 8 for calling the Trainer obj's underlying member method code with example usage: obj.train(buffer_dict) without return types