Platform For AI:Optimalkan RetinaNet dengan Blade dan plugin TensorRT

Batasan

Prosedur

Langkah 1: Buat model PyTorch dengan Plugin TensorRT

Blade dapat bekerja dengan mekanisme ekstensi TensorRT untuk optimisasi kolaboratif. Bagian ini menjelaskan cara menggunakan ekstensi TensorRT guna mengimplementasikan bagian post-processing RetinaNet. Untuk informasi lebih lanjut tentang pengembangan dan kompilasi Plugin TensorRT, lihat Dokumentasi NVIDIA Deep Learning TensorRT. Logika program untuk bagian post-processing RetinaNet dalam topik ini berasal dari komunitas open source NVIDIA. Untuk informasi lebih lanjut, lihat Retinanet-Examples. Topik ini mengekstraksi kode inti untuk menjelaskan proses pengembangan dan implementasi Custom Operator.

1. Unduh dan ekstrak kode contoh.

   wget -nv https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/tutorials/retinanet_example/retinanet-examples.tar.gz -O retinanet-examples.tar.gz
   tar xvfz retinanet-examples.tar.gz 1>/dev/null

2. Kompilasi Plugin TensorRT.

   Kode contoh mencakup implementasi dan registrasi Plugin TensorRT untuk post-processing decode dan nms RetinaNet. Dokumentasi resmi PyTorch menyediakan tiga cara untuk mengompilasi Custom Operators: Membangun dengan CMake, kompilasi Just-in-Time (JIT), dan Membangun dengan Setuptools. Untuk informasi lebih lanjut, lihat MEMPERLUAS TORCHSCRIPT DENGAN CUSTOM C++ OPERATORS. Ketiga metode kompilasi ini cocok untuk skenario berbeda; Anda dapat memilih salah satu sesuai kebutuhan. Untuk kesederhanaan, topik ini menggunakan metode kompilasi JIT. Berikut adalah kode contohnya.

   Catatan

   Sebelum dikompilasi, Anda harus mengonfigurasi pustaka dependensi seperti TensorRT, CUDA, dan CuDNN.

   import torch.utils.cpp_extension
   import os
   
   codebase="retinanet-examples"
   sources=['csrc/plugins/plugin.cpp',
            'csrc/cuda/decode.cu',
            'csrc/cuda/nms.cu',]
   sources = [os.path.join(codebase,src) for src in sources]
   torch.utils.cpp_extension.load(
       name="plugin",
       sources=sources,
       build_directory=codebase,
       extra_include_paths=['/usr/local/TensorRT/include/', '/usr/local/cuda/include/', '/usr/local/cuda/include/thrust/system/cuda/detail'],
       extra_cflags=['-std=c++14', '-O2', '-Wall'],
       extra_ldflags=['-L/usr/local/TensorRT/lib/', '-lnvinfer'],
       extra_cuda_cflags=[
           '-std=c++14', '--expt-extended-lambda',
           '--use_fast_math', '-Xcompiler', '-Wall,-fno-gnu-unique',
           '-gencode=arch=compute_75,code=sm_75',],
       is_python_module=False,
       with_cuda=True,
       verbose=False,
   )

3. Enkapsulasi bagian model konvolusi RetinaNet.

   Enkapsulasi bagian model RetinaNet ke dalam modul RetinaNetBackboneAndHeads yang terpisah.

   import torch
   from typing import List
   from torch import Tensor
   from torch.testing import assert_allclose
   from detectron2 import model_zoo
   
   # Kelas ini mengenkapsulasi backbone dan RPN heads RetinaNet.
   class RetinaNetBackboneAndHeads(torch.nn.Module):
   
       def __init__(self, model):
           super().__init__()
           self.model = model
   
       def preprocess(self, img):
           batched_inputs = [{"image": img}]
           images = self.model.preprocess_image(batched_inputs)
           return images.tensor
   
       def forward(self, images):
           features = self.model.backbone(images)
           features = [features[f] for f in self.model.head_in_features]
           cls_heads, box_heads = self.model.head(features)
           cls_heads = [cls.sigmoid() for cls in cls_heads]
           box_heads = [b.contiguous() for b in box_heads]
           return cls_heads, box_heads
   
   retinanet_model = model_zoo.get("COCO-Detection/retinanet_R_50_FPN_3x.yaml", trained=True).eval()
   retinanet_bacbone_heads = RetinaNetBackboneAndHeads(retinanet_model)

4. Bangun jaringan post-processing RetinaNet menggunakan Plugin TensorRT. Jika Anda telah membuat TensorRT Engine, Anda dapat melewati langkah ini.

   1. Buat TensorRT Engine.

      Agar Plugin TensorRT berfungsi, Anda harus mengimplementasikan fitur-fitur berikut:

      • Memuat file plugin.so yang telah dikompilasi secara dinamis menggunakan ctypes.cdll.LoadLibrary.

      • Fungsi build_retinanet_decode menggunakan API Python tensorrt untuk membangun jaringan post-processing, lalu mengompilasinya menjadi Engine.

      Berikut adalah kode contohnya.

      import os
      import numpy as np
      import tensorrt as trt
      
      import ctypes
      # Muat pustaka dinamis Plugin TensorRT.
      codebase="retinanet-examples"
      ctypes.cdll.LoadLibrary(os.path.join(codebase, 'plugin.so'))
      
      TRT_LOGGER = trt.Logger()
      trt.init_libnvinfer_plugins(TRT_LOGGER, "")
      PLUGIN_CREATORS = trt.get_plugin_registry().plugin_creator_list
      
      # Dapatkan fungsi Plugin TensorRT.
      def get_trt_plugin(plugin_name, field_collection):
          plugin = None
          for plugin_creator in PLUGIN_CREATORS:
              if plugin_creator.name != plugin_name:
                  continue
              if plugin_name == "RetinaNetDecode":
                  plugin = plugin_creator.create_plugin(
                      name=plugin_name, field_collection=field_collection
                  )
              if plugin_name == "RetinaNetNMS":
                  plugin = plugin_creator.create_plugin(
                      name=plugin_name, field_collection=field_collection
                  )
          assert plugin is not None, "plugin tidak ditemukan"
          return plugin
      
      # Fungsi untuk membangun jaringan TensorRT.
      def build_retinanet_decode(example_outputs,
              input_image_shape,
              anchors_list,
              test_score_thresh = 0.05,
              test_nms_thresh = 0.5,
              test_topk_candidates = 1000,
              max_detections_per_image = 100,
          ):
          builder = trt.Builder(TRT_LOGGER)
          EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
          network = builder.create_network(EXPLICIT_BATCH)
          config = builder.create_builder_config()
          config.max_workspace_size = 3 ** 20
      
          cls_heads, box_heads = example_outputs
          profile = builder.create_optimization_profile()
          decode_scores = []
          decode_boxes = []
          decode_class = []
      
          input_blob_names = []
          input_blob_types = []
          def _add_input(head_tensor, head_name):
              input_blob_names.append(head_name)
              input_blob_types.append("Float")
              head_shape = list(head_tensor.shape)[-3:]
              profile.set_shape(
                   head_name, [1] + head_shape, [20] + head_shape, [1000] + head_shape)
              return network.add_input(
                  name=head_name, dtype=trt.float32, shape=[-1] + head_shape
              )
      
          # Bangun input jaringan.
          cls_head_inputs = []
          cls_head_strides = [input_image_shape[-1] // cls_head.shape[-1] for cls_head in cls_heads]
          for idx, cls_head in enumerate(cls_heads):
              cls_head_name = "cls_head" + str(idx)
              cls_head_inputs.append(_add_input(cls_head, cls_head_name))
      
          box_head_inputs = []
          for idx, box_head in enumerate(box_heads):
              box_head_name = "box_head" + str(idx)
              box_head_inputs.append(_add_input(box_head, box_head_name))
      
          output_blob_names = []
          output_blob_types = []
          # Bangun jaringan decode.
          for idx, anchors in enumerate(anchors_list):
              field_coll = trt.PluginFieldCollection([
                  trt.PluginField("topk_candidates", np.array([test_topk_candidates], dtype=np.int32), trt.PluginFieldType.INT32),
                  trt.PluginField("score_thresh", np.array([test_score_thresh], dtype=np.float32), trt.PluginFieldType.FLOAT32),
                  trt.PluginField("stride", np.array([cls_head_strides[idx]], dtype=np.int32), trt.PluginFieldType.INT32),
                  trt.PluginField("num_anchors", np.array([anchors.numel()], dtype=np.int32), trt.PluginFieldType.INT32),
                  trt.PluginField("anchors", anchors.contiguous().cpu().numpy().astype(np.float32), trt.PluginFieldType.FLOAT32),]
              )
              decode_layer = network.add_plugin_v2(
                  inputs=[cls_head_inputs[idx], box_head_inputs[idx]],
                  plugin=get_trt_plugin("RetinaNetDecode", field_coll),
              )
              decode_scores.append(decode_layer.get_output(0))
              decode_boxes.append(decode_layer.get_output(1))
              decode_class.append(decode_layer.get_output(2))
      
          # Bangun jaringan NMS.
          scores_layer = network.add_concatenation(decode_scores)
          boxes_layer = network.add_concatenation(decode_boxes)
          class_layer = network.add_concatenation(decode_class)
          field_coll = trt.PluginFieldCollection([
                  trt.PluginField("nms_thresh", np.array([test_nms_thresh], dtype=np.float32), trt.PluginFieldType.FLOAT32),
                  trt.PluginField("max_detections_per_image", np.array([max_detections_per_image], dtype=np.int32), trt.PluginFieldType.INT32),]
              )
          nms_layer = network.add_plugin_v2(
             inputs=[scores_layer.get_output(0), boxes_layer.get_output(0), class_layer.get_output(0)],
             plugin=get_trt_plugin("RetinaNetNMS", field_coll),
          )
          nms_layer.get_output(0).name = "scores"
          nms_layer.get_output(1).name = "boxes"
          nms_layer.get_output(2).name = "classes"
          nms_outputs = [network.mark_output(nms_layer.get_output(k)) for k in range(3)]
          config.add_optimization_profile(profile)
          cuda_engine = builder.build_engine(network, config)
          assert cuda_engine is not None
          return cuda_engine

   2. Buat cuda_engine berdasarkan jumlah aktual output, tipe output, dan bentuk output dari RetinaNetBackboneAndHeads.

      import numpy as np
      from detectron2.data.detection_utils import read_image
      
      !wget http://images.cocodataset.org/val2017/000000439715.jpg -q -O input.jpg
      img = read_image('./input.jpg')
      img = torch.from_numpy(np.ascontiguousarray(img.transpose(2, 0, 1)))
      
      example_inputs = retinanet_bacbone_heads.preprocess(img)
      example_outputs = retinanet_bacbone_heads(example_inputs)
      
      cell_anchors = [c.contiguous() for c in retinanet_model.anchor_generator.cell_anchors]
      cuda_engine = build_retinanet_decode(
                  example_outputs, example_inputs.shape, cell_anchors)

5. Gunakan ekstensi Blade untuk mendukung model yang menggunakan PyTorch dan TensorRT Engine sekaligus.

   Kode berikut menggabungkan kembali bagian Backbone, Heads, dan post-processing Plugin TensorRT menggunakan RetinaNetWrapper, RetinaNetBackboneAndHeads, dan RetinaNetPostProcess.

   import blade.torch
   
   # Bagian post-processing yang didukung oleh ekstensi TensorRT Blade.
   class RetinaNetPostProcess(torch.nn.Module):
       def __init__(self, cuda_engine):
           super().__init__()
           blob_names = [cuda_engine.get_binding_name(idx) for idx in range(cuda_engine.num_bindings)]
           input_blob_names = blob_names[:-3]
           input_blob_types = ["Float"] * len(input_blob_names)
           output_blob_names = blob_names[-3:]
           output_blob_types = ["Float"] * len(output_blob_names)
   
           self.trt_ext_plugin = torch.classes.torch_addons.TRTEngineExtension(
               bytes(cuda_engine.serialize()),
               (input_blob_names, output_blob_names, input_blob_types, output_blob_types),
           )
   
       def forward(self, inputs: List[Tensor]):
           return self.trt_ext_plugin.forward(inputs)
   
   # Gunakan PyTorch dan TensorRT Engine bersamaan.
   class RetinaNetWrapper(torch.nn.Module):
   
       def __init__(self, model, trt_postproc):
           super().__init__()
           self.backbone_and_heads = model
           self.trt_postproc = torch.jit.script(trt_postproc)
   
       def forward(self, images):
           cls_heads, box_heads = self.backbone_and_heads(images)
           return self.trt_postproc(cls_heads + box_heads)
   
   trt_postproc = RetinaNetPostProcess(cuda_engine)
   retinanet_mix_trt = RetinaNetWrapper(retinanet_bacbone_heads, trt_postproc)
   
   # Ekspor dan simpan sebagai file TorchScript.
   retinanet_script = torch.jit.trace(retinanet_mix_trt, (example_inputs, ), check_trace=False)
   torch.jit.save(retinanet_script, 'retinanet_script.pt')
   torch.save(example_inputs, 'example_inputs.pth')
   outputs = retinanet_script(example_inputs)

   Modul torch.nn.Module yang baru dirakit memiliki fitur-fitur berikut:

   • Menggunakan antarmuka torch.classes.torch_addons.TRTEngineExtension, yang didukung oleh ekstensi TensorRT Blade.

   • Mendukung ekspor model TorchScript. Kode di atas menggunakan torch.jit.trace untuk mengekspor model.

   • Mendukung penyimpanan model dalam format TorchScript.

Langkah 2: Panggil Blade untuk mengoptimalkan model

1. Panggil antarmuka optimisasi Blade.

   Panggil antarmuka blade.optimize untuk mengoptimalkan model. Berikut adalah kode contohnya. Untuk informasi lebih lanjut tentang antarmuka blade.optimize, lihat Optimalkan model PyTorch.

   import blade
   import blade.torch
   import ctypes
   import torch
   import os
   
   codebase="retinanet-examples"
   ctypes.cdll.LoadLibrary(os.path.join(codebase, 'plugin.so'))
   
   blade_config = blade.Config()
   blade_config.gpu_config.disable_fp16_accuracy_check = True
   
   script_model = torch.jit.load('retinanet_script.pt')
   example_inputs = torch.load('example_inputs.pth')
   test_data = [(example_inputs,)] # Data input untuk PyTorch berupa daftar tupel.
   with blade_config:
       optimized_model, opt_spec, report = blade.optimize(
           script_model,  # Model TorchScript yang diekspor pada langkah sebelumnya.
           'o1',  # Aktifkan optimisasi tingkat O1 Blade.
           device_type='gpu',  # Perangkat target adalah GPU.
           test_data=test_data,  # Sediakan satu set data uji untuk membantu optimisasi dan pengujian.
       )

2. Cetak laporan optimisasi dan simpan model.

   Model yang dioptimalkan oleh Blade tetap berupa model TorchScript. Setelah optimisasi, Anda dapat menggunakan kode berikut untuk mencetak laporan optimisasi dan menyimpan model yang telah dioptimalkan.

   # Cetak laporan optimisasi.
   print("Report: {}".format(report))
   # Simpan model yang telah dioptimalkan.
   torch.jit.save(optimized_model, 'optimized.pt')

   Laporan optimisasi yang dicetak adalah sebagai berikut. Untuk informasi lebih lanjut tentang bidang-bidang dalam laporan optimisasi, lihat Laporan optimisasi.

   Report: {
     "software_context": [
       {
         "software": "pytorch",
         "version": "1.8.1+cu102"
       },
       {
         "software": "cuda",
         "version": "10.2.0"
       }
     ],
     "hardware_context": {
       "device_type": "gpu",
       "microarchitecture": "T4"
     },
     "user_config": "",
     "diagnosis": {
       "model": "unnamed.pt",
       "test_data_source": "user provided",
       "shape_variation": "undefined",
       "message": "Unable to deduce model inputs information (data type, shape, value range, etc.)",
       "test_data_info": "0 shape: (1, 3, 480, 640) data type: float32"
     },
     "optimizations": [
       {
         "name": "PtTrtPassFp16",
         "status": "effective",
         "speedup": "4.37",
         "pre_run": "40.59 ms",
         "post_run": "9.28 ms"
       }
     ],
     "overall": {
       "baseline": "40.02 ms",
       "optimized": "9.27 ms",
       "speedup": "4.32"
     },
     "model_info": {
       "input_format": "torch_script"
     },
     "compatibility_list": [
       {
         "device_type": "gpu",
         "microarchitecture": "T4"
       }
     ],
     "model_sdk": {}
   }

3. Lakukan pengujian performa pada model asli dan model yang telah dioptimalkan.

   Berikut adalah kode contoh untuk pengujian performa.

   import time
   
   @torch.no_grad()
   def benchmark(model, inp):
       for i in range(100):
           model(inp)
       torch.cuda.synchronize()
       start = time.time()
       for i in range(200):
           model(inp)
       torch.cuda.synchronize()
       elapsed_ms = (time.time() - start) * 1000
       print("Latency: {:.2f}".format(elapsed_ms / 200))
   
   # Uji performa model asli.
   benchmark(script_model, example_inputs)
   # Uji performa model yang telah dioptimalkan.
   benchmark(optimized_model, example_inputs)

   Berikut adalah hasil referensi pengujian ini.

   Latency: 40.71
   Latency: 9.35

   Hasil menunjukkan bahwa setelah 200 kali eksekusi, latensi rata-rata model asli dan model yang telah dioptimalkan masing-masing adalah 40,71 ms dan 9,35 ms.

Langkah 3: Muat dan jalankan model yang telah dioptimalkan

1. Opsional: Selama masa percobaan, tambahkan pengaturan variabel lingkungan berikut untuk mencegah program berhenti mendadak akibat kegagalan otentikasi:

   export BLADE_AUTH_USE_COUNTING=1

2. Lakukan otentikasi untuk menggunakan PAI-Blade.

   export BLADE_REGION=<region>
   export BLADE_TOKEN=<token>

   Konfigurasikan parameter berikut sesuai kebutuhan bisnis Anda:

   • <region>: wilayah tempat Anda menggunakan PAI-Blade. Anda dapat bergabung dengan grup DingTalk pengguna PAI-Blade untuk memperoleh daftar wilayah tempat PAI-Blade tersedia. Informasi kode QR grup DingTalk tersedia di Instal PAI-Blade.

   • <token>: token otentikasi yang diperlukan untuk menggunakan PAI-Blade. Anda dapat bergabung dengan grup DingTalk pengguna PAI-Blade untuk memperoleh token otentikasi tersebut. Informasi kode QR grup DingTalk tersedia di Instal PAI-Blade.

3. Muat dan jalankan model yang telah dioptimalkan.

   Model yang dioptimalkan oleh Blade tetap berupa model TorchScript. Anda dapat memuat model yang telah dioptimalkan tanpa perlu mengganti lingkungan.

   import blade.runtime.torch
   import torch
   
   from torch.testing import assert_allclose
   import ctypes
   import os
   
   codebase="retinanet-examples"
   ctypes.cdll.LoadLibrary(os.path.join(codebase, 'plugin.so'))
   
   optimized_model = torch.jit.load('optimized.pt')
   example_inputs = torch.load('example_inputs.pth')
   
   with torch.no_grad():
       pred = optimized_model(example_inputs)