text: import logging
  import os.path as osp
  
  import torch
  from torch import nn
  
  from tensorboardX import SummaryWriter
  
  from lib.test import test
  from lib.utils import checkpoint, precision_at_one, \
      Timer, AverageMeter, get_prediction, get_torch_device
  from lib.solvers import initialize_optimizer, initialize_scheduler
  
  from MinkowskiEngine import SparseTensor
  
  
  def validate(model, val_data_loader, writer, curr_iter, config, transform_data_fn):
    v_loss, v_score, v_mAP, v_mIoU = test(model, val_data_loader, config, transform_data_fn)
    writer.add_scalar('validation/mIoU', v_mIoU, curr_iter)
    writer.add_scalar('validation/loss', v_loss, curr_iter)
    writer.add_scalar('validation/precision_at_1', v_score, curr_iter)
  
    return v_mIoU
  
  
  def train(model, data_loader, val_data_loader, config, train_sampler=None, transform_data_fn=None, ngpus_per_node=1):
    device = get_torch_device(config.is_cuda)
    # Set up the train flag for batch normalization
    model.train()
  
    # Configuration
    writer = SummaryWriter(log_dir=config.log_dir)
    data_timer, iter_timer = Timer(), Timer()
    data_time_avg, iter_time_avg = AverageMeter(), AverageMeter()
    losses, scores = AverageMeter(), AverageMeter()
  
    optimizer = initialize_optimizer(model.parameters(), config)
    scheduler = initialize_scheduler(optimizer, config)
    criterion = nn.CrossEntropyLoss(ignore_index=config.ignore_label)
  
    writer = SummaryWriter(log_dir=config.log_dir)
  
    # Train the network
    logging.info('===> Start training')
    best_val_miou, best_val_iter, curr_iter, epoch, is_training = 0, 0, 1, 1, True
  
    if config.resume:
      checkpoint_fn = config.resume + '/weights.pth'
      if osp.isfile(checkpoint_fn):
        logging.info("=> loading checkpoint '{}'".format(checkpoint_fn))
        state = torch.load(checkpoint_fn)
        curr_iter = state['iteration'] + 1
        epoch = state['epoch']
        model.load_state_dict(state['state_dict'])
        if config.resume_optimizer:
          scheduler = initialize_scheduler(optimizer, config, last_step=curr_iter)
          optimizer.load_state_dict(state['optimizer'])
        if 'best_val' in state:
          best_val_miou = state['best_val']
          best_val_iter = state['best_val_iter']
        logging.info("=> loaded checkpoint '{}' (epoch {})".format(checkpoint_fn, state['epoch']))
      else:
        raise ValueError("=> no checkpoint found at '{}'".format(checkpoint_fn))
  
    while is_training:
      data_iter = data_loader.__iter__()
      if train_sampler is not None:
        train_sampler.set_epoch(epoch)
      for iteration in range(len(data_loader) // config.iter_size):
        optimizer.zero_grad()
        data_time, batch_loss = 0, 0
        iter_timer.tic()
  
        for sub_iter in range(config.iter_size):
          # Get training data
          data_timer.tic()
          coords, input, target = data_iter.next()
  
          # For some networks, making the network invariant to even, odd coords is important. Random translation
          coords[:, 1:] += (torch.rand(3) * 100).type_as(coords)
  
          # Preprocess input
          color = input[:, :3].int()
          if config.normalize_color:
            input[:, :3] = input[:, :3] / 255. - 0.5
          sinput = SparseTensor(input, coords).to(device)
  
          data_time += data_timer.toc(False)
  
          # Feed forward
          inputs = (sinput,) if config.wrapper_type == 'None' else (sinput, coords, color)
          # model.initialize_coords(*init_args)
          soutput = model(*inputs)
          # The output of the network is not sorted
          target = target.long().to(device)
  
          loss = criterion(soutput.F, target.long())
  
          # Compute and accumulate gradient
          loss /= config.iter_size
          batch_loss += loss.item()
          loss.backward()
  
        # Update number of steps
        optimizer.step()
        scheduler.step()
  
        data_time_avg.update(data_time)
        iter_time_avg.update(iter_timer.toc(False))
  
        pred = get_prediction(data_loader.dataset, soutput.F, target)
        score = precision_at_one(pred, target)
        losses.update(batch_loss, target.size(0))
        scores.update(score, target.size(0))
  
        if curr_iter >= config.max_iter:
          is_training = False
          break
  
        if curr_iter % config.stat_freq == 0 or curr_iter == 1:
          lrs = ', '.join(['{:.3e}'.format(x) for x in scheduler.get_lr()])
          debug_str = "===> Epoch[{}]({}/{}): Loss {:.4f}\tLR: {}\t".format(
              epoch, curr_iter,
              len(data_loader) // config.iter_size, losses.avg, lrs)
          debug_str += "Score {:.3f}\tData time: {:.4f}, Total iter time: {:.4f}".format(
              scores.avg, data_time_avg.avg, iter_time_avg.avg)
          logging.info(debug_str)
          # Reset timers
          data_time_avg.reset()
          iter_time_avg.reset()
          # Write logs
          writer.add_scalar('training/loss', losses.avg, curr_iter)
          writer.add_scalar('training/precision_at_1', scores.avg, curr_iter)
          writer.add_scalar('training/learning_rate', scheduler.get_lr()[0], curr_iter)
          losses.reset()
          scores.reset()
  
        # Save current status, save before val to prevent occational mem overflow
        if (curr_iter % config.save_freq == 0) and (not config.multiprocessing_distributed or (config.multiprocessing_distributed and config.rank % ngpus_per_node == 0)):
          checkpoint(model, optimizer, epoch, curr_iter, config, best_val_miou, best_val_iter)
  
        # Validation
        if (curr_iter % config.val_freq == 0) and (not config.multiprocessing_distributed or (config.multiprocessing_distributed and config.rank % ngpus_per_node == 0)):
          val_miou = validate(model, val_data_loader, writer, curr_iter, config, transform_data_fn)
          if val_miou > best_val_miou:
            best_val_miou = val_miou
            best_val_iter = curr_iter
            checkpoint(model, optimizer, epoch, curr_iter, config, best_val_miou, best_val_iter,
                       "best_val")
          logging.info("Current best mIoU: {:.3f} at iter {}".format(best_val_miou, best_val_iter))
  
          # Recover back
          model.train()
  
        # End of iteration
        curr_iter += 1
  
      epoch += 1
  
    # Explicit memory cleanup
    if hasattr(data_iter, 'cleanup'):
      data_iter.cleanup()
  
    # Save the final model
    checkpoint(model, optimizer, epoch, curr_iter, config, best_val_miou, best_val_iter)
    val_miou = validate(model, val_data_loader, writer, curr_iter, config, transform_data_fn)
    if val_miou > best_val_miou:
      best_val_miou = val_miou
      best_val_iter = curr_iter
      checkpoint(model, optimizer, epoch, curr_iter, config, best_val_miou, best_val_iter, "best_val")
    logging.info("Current best mIoU: {:.3f} at iter {}".format(best_val_miou, best_val_iter))
  