Function TEST

src/tensorrt/test/tensorrt.cpp:20–77

using namespace opr;

TEST(TestOprTensorRT, Basic) {
    REQUIRE_GPU(1);
    intl::SimpleTensorRTNetwork net;

    auto p = net.create_trt_network(true);
    auto trt_net = TensorRTOpr::to_shared_ptr_network(p.second);
    auto y2 = TensorRTOpr::make(
            TensorRTOpr::to_shared_ptr_builder(p.first), trt_net,
            intl::TensorRTGraphFeatureBits::NCHW_FLOAT, {}, {net.x})[0];

    HostTensorND host_z1;
    HostTensorND host_z2;
    auto func = net.graph->compile(
            {make_callback_copy(net.y, host_z1), make_callback_copy(y2, host_z2)});
    func->execute();
    MGB_ASSERT_TENSOR_NEAR(host_z1, host_z2, 2e-4);

    auto&& host_x = net.host_x;

// Fixme: support dynamic shape
#if NV_TENSOR_RT_VERSION < 8500
    auto&& gen = net.gen;
    *host_x = *gen({1, 23, 43, 43});
    func->execute();
    MGB_ASSERT_TENSOR_NEAR(host_z1, host_z2, 2e-4);
    *host_x = *gen({10, 23, 12, 12});
    func->execute();
    MGB_ASSERT_TENSOR_NEAR(host_z1, host_z2, 1e-3);

    *host_x = *gen({10, 23, 12, 12});
    func->execute();
    MGB_ASSERT_TENSOR_NEAR(host_z1, host_z2, 1e-3);
#endif

    // write to file so in python test we can have an engine file
    TensorRTUniquePtr<IBuilder> builder{
            createInferBuilder(TensorRTOpr::Logger::instance()), {}};
    builder->setMaxBatchSize(10);

#if NV_TENSOR_RT_VERSION >= 6001
    TensorRTUniquePtr<IBuilderConfig> build_config{builder->createBuilderConfig()};
    TensorRTUniquePtr<ICudaEngine> cuda_engine{
            builder->buildEngineWithConfig(*trt_net, *build_config)};
#else
    TensorRTUniquePtr<ICudaEngine> cuda_engine{builder->buildCudaEngine(*trt_net)};
#endif
    TensorRTUniquePtr<IHostMemory> mem{cuda_engine->serialize(), {}};
    FILE* fout = fopen(output_file("trt_cuda_engine_test").c_str(), "wb");
    auto wr = fwrite(mem->data(), 1, mem->size(), fout);
    mgb_assert(wr == mem->size());
    fclose(fout);
    debug::write_to_file(
            output_file("trt_cuda_engine_test.input").c_str(),
            debug::dump_tensor(*host_x, "x"));
    debug::write_to_file(
            output_file("trt_cuda_engine_test.output").c_str(),
            debug::dump_tensor(host_z1, "x"));
}