code: def run(
      global_rank,
      world_size,
      local_rank,
      max_epoch,
      batch_size,
      model,
      data,
      sgd,
      graph,
      verbosity,
      dist_option="plain",
      spars=None,
      precision="float32",
  ):
      dev = device.create_cuda_gpu_on(
          local_rank
      )  # need to change to CPU device for CPU-only machines
      dev.SetRandSeed(0)
      np.random.seed(0)
  
      if data == "cifar10":
          from data import cifar10
  
          train_x, train_y, val_x, val_y = cifar10.load()
      elif data == "cifar100":
          from data import cifar100
  
          train_x, train_y, val_x, val_y = cifar100.load()
      elif data == "mnist":
          from data import mnist
  
          train_x, train_y, val_x, val_y = mnist.load()
  
      num_channels = train_x.shape[1]
      image_size = train_x.shape[2]
      data_size = np.prod(train_x.shape[1 : train_x.ndim]).item()
      num_classes = (np.max(train_y) + 1).item()
  
      if model == "resnet":
          from model import resnet
  
          model = resnet.resnet50(
              num_channels=num_channels, num_classes=num_classes
          )
      elif model == "xceptionnet":
          from model import xceptionnet
  
          model = xceptionnet.create_model(
              num_channels=num_channels, num_classes=num_classes
          )
      elif model == "cnn":
          from model import cnn
  
          model = cnn.create_model(
              num_channels=num_channels, num_classes=num_classes
          )
      elif model == "alexnet":
          from model import alexnet
  
          model = alexnet.create_model(
              num_channels=num_channels, num_classes=num_classes
          )
      elif model == "mlp":
          import os, sys, inspect
  
          current = os.path.dirname(
              os.path.abspath(inspect.getfile(inspect.currentframe()))
          )
          parent = os.path.dirname(current)
          sys.path.insert(0, parent)
          from mlp import model
  
          model = model.create_model(
              data_size=data_size, num_classes=num_classes
          )
  
      # For distributed training, sequential has better performance
      if hasattr(sgd, "communicator"):
          DIST = True
          sequential = True
      else:
          DIST = False
          sequential = False
  
      if DIST:
          train_x, train_y, val_x, val_y = partition(
              global_rank, world_size, train_x, train_y, val_x, val_y
          )
  
      if model.dimension == 4:
          tx = tensor.Tensor(
              (
                  batch_size,
                  num_channels,
                  model.input_size,
                  model.input_size,
              ),
              dev,
              singa_dtype[precision],
          )
      elif model.dimension == 2:
          tx = tensor.Tensor(
              (batch_size, data_size), dev, singa_dtype[precision]
          )
          np.reshape(train_x, (train_x.shape[0], -1))
          np.reshape(val_x, (val_x.shape[0], -1))
  
      ty = tensor.Tensor((batch_size,), dev, tensor.int32)
      num_train_batch = train_x.shape[0] // batch_size
      idx = np.arange(train_x.shape[0], dtype=np.int32)
  
      # Attach model to graph
      model.set_optimizer(sgd)
      model.compile(
          [tx], is_train=True, use_graph=graph, sequential=sequential
      )
      dev.SetVerbosity(verbosity)
      model.train()
  
      # Augmentation is done only once before training
      b = 0
      x = train_x[idx[b * batch_size : (b + 1) * batch_size]]
      if model.dimension == 4:
          x = augmentation(x, batch_size)
          if image_size != model.input_size:
              x = resize_dataset(x, model.input_size)
      x = x.astype(np_dtype[precision])
      y = train_y[idx[b * batch_size : (b + 1) * batch_size]]
  
      # Copy the patch data into input tensors
      tx.copy_from_numpy(x)
      ty.copy_from_numpy(y)
  
      niters = 100
  
      # check dataset shape correctness
      if global_rank == 0:
          print("Check the shape of dataset:")
          print(tx.shape)
          print(ty.shape)
  
      # Training and evaluation loop
      # for epoch in range(max_epoch):
      #     start_time = time.time()
      #     np.random.shuffle(idx)
      #
      #     if global_rank == 0:
      #         print('Starting Epoch %d:' % (epoch))
  
      # Training phase
      dev.Sync()
      start = time.time()
  
      for b in range(niters):
          # if b % 100 == 0:
          #     print ("b: \n", b)
          # Generate the patch data in this iteration
          # Train the model
          model(tx, ty, dist_option, spars)
  
      dev.Sync()
      end = time.time()
      titer = (end - start) / float(niters)
      throughput = float(niters * batch_size * world_size) / (
          end - start
      )
  
      if global_rank == 0:
          print(
              "Throughput = {} per second".format(throughput),
              flush=True,
          )
          print("TotalTime={}".format(end - start), flush=True)
          print("Total={}".format(titer), flush=True)
          print("world_size={}".format(world_size), flush=True)
          print("batch_size={}".format(batch_size), flush=True)
          print(
              "model.input_size={}".format(model.input_size), flush=True
          )
          print("num_channels={}".format(num_channels), flush=True)
          print("num_classes={}".format(num_classes), flush=True)
          print("data_size={}".format(data_size), flush=True)
          print("image_size={}".format(image_size), flush=True)
  
      dev.PrintTimeProfiling()
  
file_path: singa/examples/cifar_distributed_cnn/train_cnn.py
imports: from singa import singa_wrap
  from singa import device
  from singa import tensor
  from singa import opt
  import numpy
  import time
  import argparse
  from PIL import Image
  import warnings
label_desc: Use image node_id 6 for calling a global function with example usage: run(global_rank, world_size, local_rank, max_epoch, batch_size, model, data, sgd, graph, verbosity, dist_option, spars, precision) without return types