code: def train_one_epoch(
      model,
      optimizer,
      data_loader,
      device,
      epoch,
      num_classes,
      lr_scheduler,
      print_freq=10,
      scaler=None,
  ):
      model.train()
      metric_logger = utils.MetricLogger(delimiter="  ")
      metric_logger.add_meter(
          "lr", utils.SmoothedValue(window_size=1, fmt="{value:.6f}")
      )
      header = "Epoch: [{}]".format(epoch)
  
      if num_classes == 2:
          # 设置cross_entropy中背景和前景的loss权重(根据自己的数据集进行设置)
          loss_weight = torch.as_tensor([1.0, 2.0], device=device)
      else:
          loss_weight = None
  
      for image, target in metric_logger.log_every(
          data_loader, print_freq, header
      ):
          image, target = image.to(device), target.to(device)
          with torch.cuda.amp.autocast(enabled=scaler is not None):
              output = model(image)
              loss = criterion(
                  output,
                  target,
                  loss_weight,
                  num_classes=num_classes,
                  ignore_index=255,
              )
  
          optimizer.zero_grad()
          if scaler is not None:
              scaler.scale(loss).backward()
              scaler.step(optimizer)
              scaler.update()
          else:
              loss.backward()
              optimizer.step()
  
          lr_scheduler.step()
  
          lr = optimizer.param_groups[0]["lr"]
          metric_logger.update(loss=loss.item(), lr=lr)
  
      return metric_logger.meters["loss"].global_avg, lr
  
file_path: deep-learning-for-image-processing/pytorch_segmentation/unet/train_utils/train_and_eval.py
imports: import torch
  from torch import nn
  import train_utils.distributed_utils
  from .dice_coefficient_loss import dice_loss, build_target
label_desc: Use image node_id 3 for calling a global function with example usage: train_one_epoch(model, optimizer, data_loader, device, epoch, num_classes, lr_scheduler, print_freq, scaler) and returns: metric_logger,  lr