code: def main():
      parser = argparse.ArgumentParser()
      parser.add_argument("--data-dir", default="auto")
      parser.add_argument("--dataset", choices=("ade20k", "cityscapes"))
      parser.add_argument(
          "--model", choices=("pspnet_resnet101", "pspnet_resnet50")
      )
      parser.add_argument("--lr", default=1e-2, type=float)
      parser.add_argument("--batchsize", default=2, type=int)
      parser.add_argument("--out", default="result")
      parser.add_argument("--iteration", default=None, type=int)
      parser.add_argument("--communicator", default="pure_nccl")
      args = parser.parse_args()
  
      dataset_cfgs = {
          "ade20k": {
              "input_size": (473, 473),
              "label_names": ade20k_semantic_segmentation_label_names,
              "iteration": 150000,
          },
          "cityscapes": {
              "input_size": (713, 713),
              "label_names": cityscapes_semantic_segmentation_label_names,
              "iteration": 90000,
          },
      }
      dataset_cfg = dataset_cfgs[args.dataset]
  
      # https://docs.chainer.org/en/stable/chainermn/tutorial/tips_faqs.html#using-multiprocessiterator
      if hasattr(multiprocessing, "set_start_method"):
          multiprocessing.set_start_method("forkserver")
          p = multiprocessing.Process()
          p.start()
          p.join()
  
      comm = chainermn.create_communicator(args.communicator)
      device = comm.intra_rank
  
      n_class = len(dataset_cfg["label_names"])
      if args.model == "pspnet_resnet101":
          model = PSPNetResNet101(
              n_class,
              pretrained_model="imagenet",
              input_size=dataset_cfg["input_size"],
          )
      elif args.model == "pspnet_resnet50":
          model = PSPNetResNet50(
              n_class,
              pretrained_model="imagenet",
              input_size=dataset_cfg["input_size"],
          )
      train_chain = create_mnbn_model(TrainChain(model), comm)
      model = train_chain.model
      if device >= 0:
          chainer.cuda.get_device_from_id(device).use()
          train_chain.to_gpu()
  
      if args.iteration is None:
          n_iter = dataset_cfg["iteration"]
      else:
          n_iter = args.iteration
  
      if args.dataset == "ade20k":
          train = ADE20KSemanticSegmentationDataset(
              data_dir=args.data_dir, split="train"
          )
          if comm.rank == 0:
              val = ADE20KSemanticSegmentationDataset(
                  data_dir=args.data_dir, split="val"
              )
          label_names = ade20k_semantic_segmentation_label_names
      elif args.dataset == "cityscapes":
          train = CityscapesSemanticSegmentationDataset(
              args.data_dir, label_resolution="fine", split="train"
          )
          if comm.rank == 0:
              val = CityscapesSemanticSegmentationDataset(
                  args.data_dir, label_resolution="fine", split="val"
              )
          label_names = cityscapes_semantic_segmentation_label_names
      train = TransformDataset(
          train,
          ("img", "label"),
          Transform(model.mean, dataset_cfg["input_size"]),
      )
  
      if comm.rank == 0:
          indices = np.arange(len(train))
      else:
          indices = None
      indices = chainermn.scatter_dataset(indices, comm, shuffle=True)
      train = train.slice[indices]
  
      train_iter = chainer.iterators.MultiprocessIterator(
          train, batch_size=args.batchsize, n_processes=2
      )
  
      optimizer = chainermn.create_multi_node_optimizer(
          chainer.optimizers.MomentumSGD(args.lr, 0.9), comm
      )
      optimizer.setup(train_chain)
      for param in train_chain.params():
          if param.name not in ("beta", "gamma"):
              param.update_rule.add_hook(
                  chainer.optimizer.WeightDecay(1e-4)
              )
      for l in [
          model.ppm,
          model.head_conv1,
          model.head_conv2,
          train_chain.aux_conv1,
          train_chain.aux_conv2,
      ]:
          for param in l.params():
              param.update_rule.add_hook(GradientScaling(10))
  
      updater = training.updaters.StandardUpdater(
          train_iter, optimizer, device=device
      )
      trainer = training.Trainer(
          updater, (n_iter, "iteration"), args.out
      )
      trainer.extend(
          PolynomialShift("lr", 0.9, n_iter, optimizer=optimizer),
          trigger=(1, "iteration"),
      )
  
      log_interval = 10, "iteration"
  
      if comm.rank == 0:
          trainer.extend(extensions.LogReport(trigger=log_interval))
          trainer.extend(extensions.observe_lr(), trigger=log_interval)
          trainer.extend(
              extensions.PrintReport(
                  [
                      "epoch",
                      "iteration",
                      "elapsed_time",
                      "lr",
                      "main/loss",
                      "validation/main/miou",
                      "validation/main/mean_class_accuracy",
                      "validation/main/pixel_accuracy",
                  ]
              ),
              trigger=log_interval,
          )
          trainer.extend(extensions.ProgressBar(update_interval=10))
          trainer.extend(
              extensions.snapshot_object(
                  train_chain.model,
                  "snapshot_model_{.updater.iteration}.npz",
              ),
              trigger=(n_iter, "iteration"),
          )
          val_iter = chainer.iterators.SerialIterator(
              val, batch_size=1, repeat=False, shuffle=False
          )
          trainer.extend(
              SemanticSegmentationEvaluator(
                  val_iter, model, label_names
              ),
              trigger=(n_iter, "iteration"),
          )
  
      trainer.run()
  
imports: import argparse
  import copy
  import multiprocessing
  import numpy
  import chainer
  import chainer.functions
  import chainer.links
  from chainer import training
  from chainer.training import extensions
  from chainer.training.extensions import PolynomialShift
  from chainercv.datasets import ade20k_semantic_segmentation_label_names
  from chainercv.datasets import ADE20KSemanticSegmentationDataset
  from chainercv.datasets import cityscapes_semantic_segmentation_label_names
  from chainercv.datasets import CityscapesSemanticSegmentationDataset
  from chainercv.experimental.links import PSPNetResNet101
  from chainercv.experimental.links import PSPNetResNet50
  from chainercv.chainer_experimental.datasets.sliceable import TransformDataset
  from chainercv.extensions import SemanticSegmentationEvaluator
  from chainercv.links import Conv2DBNActiv
  from chainercv import transforms
  from chainercv.links.model.ssd import GradientScaling
  import PIL
  import chainermn
label_desc: Use image node_id 2 for calling a global function with example usage: main() without return types