AWS Neo-AI DLRのソースコード解析(その2)

@Vengineerの戯言 : Twitter
SystemVerilogの世界へようこそ、すべては、SystemC v0.9公開から始まった

昨日の続く、AWS Neo-AI DLRのソースコード解析(その1)

DLRModelクラスは、ここで定義されています。

DLRModel::DLRModel(const std::string& model_path,
                   const DLContext& ctx) {
  backend_ = get_backend(model_path);
  ctx_ = ctx;

//  std::string cmdline = "tar -xf " + model_path; 
//  system(cmdline.c_str());
  if (backend_ == DLRBackend::kTVM) {
    SetupTVMModule(model_path);
  } else if (backend_ == DLRBackend::kTREELITE) {
    SetupTreeliteModule(model_path);
  } else {
    LOG(FATAL) << "Unsupported backend!";
  }
}

バックエンド(TVM or Treelite)によって、生成するモデルを変えています。

get_backendメソッドにて、バックエンドを決めます。

DLRBackend get_backend(const std::string &dirname) {
  std::vector<std::string> paths;
  listdir(dirname, paths);
  DLRBackend backend = DLRBackend::kTREELITE;
  for (auto filename: paths) {
    if (endsWith(filename, ".params")){
      backend = DLRBackend::kTVM;
      break;
    }
  }
  return backend;
}

モデルがあるディレクトリ内のファイルとして、".params"で終わるファイルがあるときは、TVMがバックエンドで
それ以外の時は、Treeliteがバックエンドになります。