text: # Licensed to the Apache Software Foundation (ASF) under one
  # or more contributor license agreements.  See the NOTICE file
  # distributed with this work for additional information
  # regarding copyright ownership.  The ASF licenses this file
  # to you under the Apache License, Version 2.0 (the
  # "License"); you may not use this file except in compliance
  # with the License.  You may obtain a copy of the License at
  #
  #   http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing,
  # software distributed under the License is distributed on an
  # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  # KIND, either express or implied.  See the License for the
  # specific language governing permissions and limitations
  # under the License.
  
  import gc
  import gluoncv
  import mxnet as mx
  import numpy as np
  
  from mxnet import gluon
  from time import time
  
  from mxnet.gluon.data.vision import transforms
  
  
  def get_classif_model(model_name, use_tensorrt, ctx=mx.gpu(0), batch_size=128):
      mx.contrib.tensorrt.set_use_fp16(False)
      h, w = 32, 32
      net = gluoncv.model_zoo.get_model(model_name, pretrained=True)
      net.hybridize()
      net.forward(mx.nd.zeros((batch_size, 3, h, w)))
      net.export(model_name)
      _sym, arg_params, aux_params = mx.model.load_checkpoint(model_name, 0)
      if use_tensorrt:
          sym = _sym.get_backend_symbol('TensorRT')
          arg_params, aux_params = mx.contrib.tensorrt.init_tensorrt_params(sym, arg_params,
                                                                            aux_params)
      else:
          sym = _sym
      executor = sym.simple_bind(ctx=ctx, data=(batch_size, 3, h, w),
                                 softmax_label=(batch_size,),
                                 grad_req='null', force_rebind=True)
      executor.copy_params_from(arg_params, aux_params)
      return executor
  
  
  def cifar10_infer(model_name, use_tensorrt, num_workers, ctx=mx.gpu(0), batch_size=128):
      executor = get_classif_model(model_name, use_tensorrt, ctx, batch_size)
  
      num_ex = 10000
      all_preds = np.zeros([num_ex, 10])
  
      all_label_test = np.zeros(num_ex)
  
      transform_test = transforms.Compose([
          transforms.ToTensor(),
          transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010])
      ])
  
      data_loader = lambda: gluon.data.DataLoader(
          gluon.data.vision.CIFAR10(train=False).transform_first(transform_test),
          batch_size=batch_size, shuffle=False, num_workers=num_workers)
  
      val_data = data_loader()
  
      for idx, (data, label) in enumerate(val_data):
          # Skip last batch if it's undersized.
          if data.shape[0] < batch_size:
              continue
          offset = idx * batch_size
          all_label_test[offset:offset + batch_size] = label.asnumpy()
  
          # warm-up, but don't use result
          executor.forward(is_train=False, data=data)
          executor.outputs[0].wait_to_read()
  
      gc.collect()
      val_data = data_loader()
      example_ct = 0
      start = time()
  
      # if use_tensorrt:
      for idx, (data, label) in enumerate(val_data):
          # Skip last batch if it's undersized.
          if data.shape[0] < batch_size:
              continue
          executor.forward(is_train=False, data=data)
          preds = executor.outputs[0].asnumpy()
          offset = idx * batch_size
          all_preds[offset:offset + batch_size, :] = preds[:batch_size]
          example_ct += batch_size
  
      all_preds = np.argmax(all_preds, axis=1)
      matches = (all_preds[:example_ct] == all_label_test[:example_ct]).sum()
      duration = time() - start
  
      return duration, 100.0 * matches / example_ct
  
  
  def run_experiment_for(model_name, batch_size, num_workers):
      print("\n===========================================")
      print("Model: %s" % model_name)
      print("===========================================")
      print("*** Running inference using pure MXNet ***\n")
      mx_duration, mx_pct = cifar10_infer(model_name=model_name, batch_size=batch_size,
                                          num_workers=num_workers, use_tensorrt=False)
      print("\nMXNet: time elapsed: %.3fs, accuracy: %.2f%%" % (mx_duration, mx_pct))
      print("\n*** Running inference using MXNet + TensorRT ***\n")
      trt_duration, trt_pct = cifar10_infer(model_name=model_name, batch_size=batch_size,
                                            num_workers=num_workers, use_tensorrt=True)
      print("TensorRT: time elapsed: %.3fs, accuracy: %.2f%%" % (trt_duration, trt_pct))
      speedup = mx_duration / trt_duration
      print("TensorRT speed-up (not counting compilation): %.2fx" % speedup)
  
      acc_diff = abs(mx_pct - trt_pct)
      print("Absolute accuracy difference: %f" % acc_diff)
      return speedup, acc_diff
  
  
  def test_tensorrt_on_cifar_resnets(batch_size=32, tolerance=0.1, num_workers=1):
      original_use_fp16 = mx.contrib.tensorrt.get_use_fp16()
      try:
          models = [
              'cifar_resnet20_v1',
              'cifar_resnet56_v1',
              'cifar_resnet110_v1',
              'cifar_resnet20_v2',
              'cifar_resnet56_v2',
              'cifar_resnet110_v2',
              'cifar_wideresnet16_10',
              'cifar_wideresnet28_10',
              'cifar_wideresnet40_8',
              'cifar_resnext29_16x64d'
          ]
  
          num_models = len(models)
  
          speedups = np.zeros(num_models, dtype=np.float32)
          acc_diffs = np.zeros(num_models, dtype=np.float32)
  
          test_start = time()
  
          for idx, model in enumerate(models):
              speedup, acc_diff = run_experiment_for(model, batch_size, num_workers)
              speedups[idx] = speedup
              acc_diffs[idx] = acc_diff
              assert acc_diff < tolerance, "Accuracy difference between MXNet and TensorRT > %.2f%% for model %s" % (
                  tolerance, model)
  
          print("Perf and correctness checks run on the following models:")
          print(models)
          mean_speedup = np.mean(speedups)
          std_speedup = np.std(speedups)
          print("\nSpeedups:")
          print(speedups)
          print("Speedup range: [%.2f, %.2f]" % (np.min(speedups), np.max(speedups)))
          print("Mean speedup: %.2f" % mean_speedup)
          print("St. dev. of speedups: %.2f" % std_speedup)
          print("\nAcc. differences: %s" % str(acc_diffs))
  
          test_duration = time() - test_start
  
          print("Test duration: %.2f seconds" % test_duration)
      finally:
          mx.contrib.tensorrt.set_use_fp16(original_use_fp16)
  
  