Function main

example/cpp/bert_model_example.cpp:124–148

using namespace turbo_transformers;

int main(int argc, char *argv[]) {
  if (argc != 2) {
    std::cerr << "./bert_example npz_model_path" << std::endl;
    return -1;
  }
  const std::string model_path = static_cast<std::string>(argv[1]);

  if (core::IsCompiledWithCUDA()) {
    std::cout << "run bert on GPU, device id is 0" << std::endl;
    test(model_path, true /*use cuda*/);
  }
  std::cout << "run bert on CPU, use 4 threads to do bert inference"
            << std::endl;
  turbo_transformers::core::SetNumThreads(4);
  test(model_path, false /*not use cuda*/);
  turbo_transformers::core::SetNumThreads(1);
  if (core::IsCompiledWithCUDA()) {
    std::cout << "10 threads do 10 independent bert inferences." << std::endl;
    test_multiple_threads(model_path, false /*only_input*/, true /*use cuda*/,
                          10);
  }
  test_multiple_threads(model_path, false /*only_input*/,
                        false /*not use cuda*/, 10);
  return 0;
}