code: import numpy as np
  
  import time
  import tempfile
  import os
  import importlib.util
  import argparse
  from typing import Sequence
  import subprocess
  import re
  
  
  import oneflow as flow
  import oneflow._oneflow_internal as oneflow_internal
  
  
  DEFAULT_TIMES = 20
  gpu_memory_used_by_oneflow = 0
  
  
  def import_file(path):
      spec = importlib.util.spec_from_file_location("mod", path)
      mod = importlib.util.module_from_spec(spec)
      spec.loader.exec_module(mod)
      return mod
  
  
  def sync(x):
      if test_oneflow:
          x.numpy()
      else:
          x.cpu()
  
  
  def gpu_memory_used():
      output = subprocess.check_output(
          [
              "nvidia-smi",
              "--query-compute-apps=pid,used_gpu_memory",
              "--format=csv,noheader",
          ]
      )
      output = output.decode("utf-8").strip()
      my_pid = os.getpid()
      mem_used_by_me = 0
      for line in output.split("\n"):
          pid, mem_used = map(int, re.split(",? ", line)[:2])
          if pid == my_pid:
              mem_used_by_me += mem_used
      return mem_used_by_me
  
  
  def print_rank_0(*args, **kwargs):
      rank = int(os.getenv("RANK", "0"))
      if rank == 0:
          print(*args, **kwargs)
  
  
  def test(
      model_path: str,
      module_name: str,
      input_shape: Sequence[int],
      disable_backward=False,
      times=DEFAULT_TIMES,
      no_verbose=False,
      ddp=False,
      ddp_broadcast_buffers=False,
      show_memory=True,
  ):
      framework_name = "OneFlow" if test_oneflow else "PyTorch"
      if test_oneflow:
          python_module = import_file(model_path)
          torch = flow
      else:
          with open(model_path) as f:
              buf = f.read()
  
          lines = buf.split("\n")
          for i, line in enumerate(lines):
              if "import" not in line and len(line.strip()) != 0:
                  break
          lines = (
              lines[:i]
              + [
                  "import torch as flow",
                  "import torch.nn as nn",
                  "from torch import Tensor",
                  "from torch.nn import Parameter",
              ]
              + lines[i:]
          )
          buf = "\n".join(lines)
          with tempfile.NamedTemporaryFile("w", suffix=".py") as f:
              f.write(buf)
              f.flush()
              python_module = import_file(f.name)
  
          import torch
  
          if ddp:
              import torch.distributed as dist
  
              local_rank_env_var = os.getenv("LOCAL_RANK")
              assert local_rank_env_var is not None
              rank = int(local_rank_env_var)
              torch.cuda.set_device(rank)
  
              dist.init_process_group(backend="nccl", init_method="env://")
  
      Net = getattr(python_module, module_name)
  
      warmup_times = 5
  
      m = Net()
      m = m.to("cuda")
  
      if ddp:
          if test_oneflow:
              m = torch.nn.parallel.DistributedDataParallel(
                  m, broadcast_buffers=ddp_broadcast_buffers
              )
          else:
              m = torch.nn.parallel.DistributedDataParallel(
                  m, device_ids=[rank], broadcast_buffers=ddp_broadcast_buffers
              )
  
      def run_model(m, x):
          if disable_backward:
              with torch.no_grad():
                  return m(x)
          else:
              return m(x)
  
      learning_rate = 0.01
      mom = 0.9
      optimizer = torch.optim.SGD(m.parameters(), lr=learning_rate, momentum=mom)
  
      # input tensor of OneFlow should set requires_grad=False due to a bug
      x = torch.tensor(
          np.ones(input_shape).astype(np.float32), requires_grad=not test_oneflow
      ).to("cuda")
      for i in range(warmup_times + times):
          if i == warmup_times:
              start = time.time()
          y = run_model(m, x)
          if not disable_backward:
              y = y.sum()
              y.backward()
              optimizer.zero_grad()
              optimizer.step()
          sync(y)
      end = time.time()
      total_time_ms = (end - start) * 1000
      time_per_run_ms = total_time_ms / times
      if no_verbose:
          print_rank_0(f"{framework_name}: {time_per_run_ms:.1f}ms")
      else:
          print_rank_0(
              f"{framework_name} {module_name} time: {time_per_run_ms:.1f}ms (= {total_time_ms:.1f}ms / {times}, input_shape={input_shape}{', backward is disabled' if disable_backward else ''}{', ddp' if ddp else ''}{', ddp_broadcast_buffers is disabled' if not ddp_broadcast_buffers else ''}{f', world size={flow.env.get_world_size()}' if flow.env.get_world_size() != 1 else ''})"
          )
      if show_memory:
          global gpu_memory_used_by_oneflow
          if test_oneflow:
              gpu_memory_used_by_oneflow = gpu_memory_used()
  
              print_rank_0(
                  f"{framework_name} GPU used (rank 0): {gpu_memory_used_by_oneflow} MiB"
              )
          else:
              print_rank_0(
                  f"{framework_name} GPU used (rank 0, estimated): {gpu_memory_used() - gpu_memory_used_by_oneflow} MiB"
              )
      if ddp and not test_oneflow:
          import torch.distributed as dist
  
          dist.destroy_process_group()
  
      return time_per_run_ms
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser()
      parser.add_argument("model_path", type=str)
      parser.add_argument("module_name", type=str)
      parser.add_argument("input_shape", type=str)
      parser.add_argument("--times", type=int, default=DEFAULT_TIMES)
      parser.add_argument("--disable-backward", action="store_true")
      parser.add_argument("--no-verbose", action="store_true")
      parser.add_argument("--ddp", action="store_true")
      parser.add_argument("--ddp-no-broadcast-buffers", action="store_true")
      parser.add_argument("--only-oneflow", action="store_true")
      parser.add_argument("--only-pytorch", action="store_true")
      parser.add_argument("--no-show-memory", action="store_true")
  
      args = parser.parse_args()
      input_shape = list(map(int, args.input_shape.split("x")))
  
      global test_oneflow
  
      if not args.only_pytorch:
          # NOTE: PyTorch must run after OneFlow for correct memory usage
          test_oneflow = True
          oneflow_time = test(
              args.model_path,
              args.module_name,
              input_shape,
              disable_backward=args.disable_backward,
              times=args.times,
              no_verbose=args.no_verbose,
              ddp=args.ddp,
              ddp_broadcast_buffers=not args.ddp_no_broadcast_buffers,
              show_memory=not args.no_show_memory,
          )
  
      if not args.only_oneflow:
          test_oneflow = False
          pytorch_time = test(
              args.model_path,
              args.module_name,
              input_shape,
              disable_backward=args.disable_backward,
              times=args.times,
              no_verbose=args.no_verbose,
              ddp=args.ddp,
              ddp_broadcast_buffers=not args.ddp_no_broadcast_buffers,
              show_memory=not args.no_show_memory,
          )
  
      if not args.only_pytorch and not args.only_oneflow:
          relative_speed = pytorch_time / oneflow_time
          if args.no_verbose:
              print_rank_0(f"Relative speed: {relative_speed:.2f}")
          else:
              print_rank_0(
                  f"Relative speed: {relative_speed:.2f} (= {pytorch_time:.1f}ms / {oneflow_time:.1f}ms)"
              )
   import oneflow as flow
  import argparse
  import numpy as np
  import os
  import time
  
  from models.posenet import PoseNet
  
  from utils.ofrecord_data_utils import OFRecordDataLoader
  
  
  def _parse_args():
      parser = argparse.ArgumentParser("flags for train posenet")
      parser.add_argument(
          "--save_checkpoint_path",
          type=str,
          default="./checkpoints",
          help="save checkpoint root dir",
      )
      parser.add_argument(
          "--load_checkpoint", type=str, default="", help="load checkpoint"
      )
      parser.add_argument(
          "--ofrecord_path", type=str, default="./ofrecord", help="dataset path"
      )
      # training hyper-parameters
      parser.add_argument(
          "--learning_rate", type=float, default=0.001, help="learning rate"
      )
      parser.add_argument("--mom", type=float, default=0.9, help="momentum")
      parser.add_argument("--epochs", type=int, default=1000, help="training epochs")
      parser.add_argument(
          "--train_batch_size", type=int, default=32, help="train batch size"
      )
      parser.add_argument("--val_batch_size", type=int, default=32, help="val batch size")
  
      return parser.parse_args()
  
  
  def main(args):
      flow.enable_eager_execution()
  
      train_data_loader = OFRecordDataLoader(
          ofrecord_root=args.ofrecord_path,
          mode="train",
          # NOTE(Liang Depeng): needs to explictly set the dataset size
          dataset_size=7459,
          batch_size=args.train_batch_size,
      )
  
      val_data_loader = OFRecordDataLoader(
          ofrecord_root=args.ofrecord_path,
          mode="val",
          dataset_size=1990,
          batch_size=args.val_batch_size,
      )
  
      # oneflow init
      start_t = time.time()
      posenet_module = PoseNet()
  
      if args.load_checkpoint != "":
          posenet_module.load_state_dict(flow.load(args.load_checkpoint))
  
      end_t = time.time()
      print("init time : {}".format(end_t - start_t))
  
      of_cross_entropy = flow.nn.CrossEntropyLoss()
  
      posenet_module.to("cuda")
      of_cross_entropy.to("cuda")
  
      of_sgd = flow.optim.SGD(
          posenet_module.parameters(), lr=args.learning_rate, momentum=args.mom
      )
  
      of_losses = []
      all_samples = len(val_data_loader) * args.val_batch_size
      print_interval = 100
  
      for epoch in range(args.epochs):
          posenet_module.train()
  
          for b in range(len(train_data_loader)):
              image, label = train_data_loader.get_batch()
  
              # oneflow train
              start_t = time.time()
              image = image.to("cuda")
              label = label.to("cuda")
              logits = posenet_module(image)
              loss = of_cross_entropy(logits, label)
  
              loss.backward()
              of_sgd.step()
              of_sgd.zero_grad()
              end_t = time.time()
              if b % print_interval == 0:
                  l = loss.numpy()
                  of_losses.append(l)
                  print(
                      "epoch {} train iter {} oneflow loss {}, train time : {}".format(
                          epoch, b, l, end_t - start_t
                      )
                  )
  
          print("epoch %d train done, start validation" % epoch)
          posenet_module.eval()
          correct_of = 0.0
          for b in range(len(val_data_loader)):
              image, label = val_data_loader.get_batch()
  
              start_t = time.time()
              image = image.to("cuda")
              with flow.no_grad():
                  logits = posenet_module(image)
                  predictions = logits.softmax()
              of_predictions = predictions.numpy()
              clsidxs = np.argmax(of_predictions, axis=1)
  
              label_nd = label.numpy()
              for i in range(args.val_batch_size):
                  if clsidxs[i] == label_nd[i]:
                      correct_of += 1
              end_t = time.time()
  
          print("epoch %d, oneflow top1 val acc: %f" % (epoch, correct_of / all_samples))
          flow.save(
              posenet_module.state_dict(),
              os.path.join(
                  args.save_checkpoint_path,
                  "epoch_%d_val_acc_%f" % (epoch, correct_of / all_samples),
              ),
          )
  
      writer = open("of_losses.txt", "w")
      for o in of_losses:
          writer.write("%f\n" % o)
      writer.close()
  
  
  if __name__ == "__main__":
      args = _parse_args()
      main(args)
  
apis: numpy.ones torch.distributed.init_process_group torch.no_grad torch.nn.parallel.DistributedDataParallel torch.distributed.destroy_process_group torch.cuda.set_device
apis: numpy.argmax