tensorrt_runtime.cpp 6.1 KB
Newer Older
1 2 3 4
/**
 * \file src/tensorrt/test/tensorrt_runtime.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
5
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */

#include "megbrain/comp_node_env.h"
#include "megbrain/test/autocheck.h"
#include "megbrain/test/helper.h"
#include "megbrain/test/megdnn_helper.h"
#include "megbrain/utils/debug.h"

#if MGB_ENABLE_TENSOR_RT

#include "megbrain/tensorrt/tensorrt_opr.h"
#include "megbrain/tensorrt/tensorrt_runtime_opr.h"
#include "make_trt_net.h"

#include <random>

using namespace mgb;
using namespace nvinfer1;

template <typename T>
using TensorRTUniquePtr = intl::TensorRTUniquePtr<T>;



TEST(TestOprTensorRT, RuntimeBasic) {
    REQUIRE_GPU(1);
    intl::SimpleTensorRTNetwork net;
    auto make_trt = [&net]() {
        auto p = net.create_trt_network(false);
        TensorRTUniquePtr<INetworkDefinition> trt_net{p.second, {}};
        TensorRTUniquePtr<IBuilder> builder{p.first, {}};
        builder->setMaxBatchSize(5);
#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(), {}};
        return TensorRTRuntimeOpr::make(mem->data(), mem->size(), {net.x})[0];
    };
    auto y2 = make_trt();

    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();
61
    MGB_ASSERT_TENSOR_NEAR(host_z1, host_z2, 5e-4);
62 63 64
}


65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95
TEST(TestOprTensorRT, RuntimeBasicBatched) {
    REQUIRE_GPU(1);
    intl::BatchedTensorRTNetwork net;
    auto make_trt = [&net]() {
        auto p = net.create_trt_network(false);
        TensorRTUniquePtr<INetworkDefinition> trt_net{p.second, {}};
        TensorRTUniquePtr<IBuilder> builder{p.first, {}};
        builder->setMaxBatchSize(5);
#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(), {}};
        auto nx = opr::Broadcast::make(net.x, {1, net.x.shape()[0], net.x.shape()[1], net.x.shape()[2]});
        return TensorRTRuntimeOpr::make(mem->data(), mem->size(), {nx})[0];
    };
    auto y2 = make_trt();

    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, 5e-4);
}

96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161

TEST(TestOprTensorRT, ConcatRuntimeBasic) {
    REQUIRE_GPU(1);
    intl::ConcatConvTensorRTNetwork net;

    SymbolVar y2;
    {
        auto p = net.create_trt_network(false);
        TensorRTUniquePtr<INetworkDefinition> trt_net{p.second, {}};
        TensorRTUniquePtr<IBuilder> builder{p.first, {}};
        builder->setMaxBatchSize(5);
#if NV_TENSOR_RT_VERSION >= 6001
        TensorRTUniquePtr<IBuilderConfig> build_config{
                builder->createBuilderConfig()};
        auto cuda_engine =
                builder->buildEngineWithConfig(*trt_net, *build_config);
#else
        auto cuda_engine = builder->buildCudaEngine(*trt_net);
#endif
        TensorRTUniquePtr<IHostMemory> mem{cuda_engine->serialize(), {}};

        FILE* fout = fopen(output_file("trt_cuda_engine").c_str(), "wb");
        auto wr = fwrite(mem->data(), 1, mem->size(), fout);
        mgb_assert(wr == mem->size());
        fclose(fout);

        y2 = TensorRTRuntimeOpr::make(
                TensorRTRuntimeOpr::to_shared_ptr_engine(cuda_engine), {},
                {net.x0, net.x1})[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, 1e-4);
}

TEST(TestOprTensorRT, RuntimeChangeBatchSize) {
    REQUIRE_GPU(1);
    intl::SimpleTensorRTNetwork net;
    auto make_trt = [&net]() {
        auto p = net.create_trt_network(false);
        TensorRTUniquePtr<INetworkDefinition> trt_net{p.second, {}};
        TensorRTUniquePtr<IBuilder> builder{p.first, {}};
        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(), {}};
        return TensorRTRuntimeOpr::make(mem->data(), mem->size(), {net.x})[0];
    };
    auto y2 = make_trt();

    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();
162
    MGB_ASSERT_TENSOR_NEAR(host_z1, host_z2, 5e-4);
163 164
    *net.host_x = *net.gen({1, 23, 28, 28});
    func->execute();
165
    MGB_ASSERT_TENSOR_NEAR(host_z1, host_z2, 5e-4);
166 167
    *net.host_x = *net.gen({10, 23, 28, 28});
    func->execute();
168
    MGB_ASSERT_TENSOR_NEAR(host_z1, host_z2, 5e-4);
169 170 171 172 173
}

#endif  // MGB_ENABLE_TENSOR_RT

// vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}