From 161e6b05ee7bf060e767c3a3f4cb1afccde6c34c Mon Sep 17 00:00:00 2001 From: JumpPandaer <834662983@qq.com> Date: Thu, 1 Jul 2021 13:09:51 +0800 Subject: [PATCH] add detr (#606) * add detr * Update README.md --- detr/CMakeLists.txt | 33 ++ detr/README.md | 58 ++++ detr/backbone.hpp | 326 +++++++++++++++++ detr/common.hpp | 90 +++++ detr/detr.cpp | 826 ++++++++++++++++++++++++++++++++++++++++++++ detr/gen_wts.py | 123 +++++++ detr/logging.h | 447 ++++++++++++++++++++++++ 7 files changed, 1903 insertions(+) create mode 100644 detr/CMakeLists.txt create mode 100644 detr/README.md create mode 100644 detr/backbone.hpp create mode 100644 detr/common.hpp create mode 100644 detr/detr.cpp create mode 100644 detr/gen_wts.py create mode 100644 detr/logging.h diff --git a/detr/CMakeLists.txt b/detr/CMakeLists.txt new file mode 100644 index 0000000..2b4e145 --- /dev/null +++ b/detr/CMakeLists.txt @@ -0,0 +1,33 @@ +cmake_minimum_required(VERSION 2.6) + +project(detr) + +add_definitions(-std=c++11) + +option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) +set(CMAKE_CXX_STANDARD 11) +set(CMAKE_BUILD_TYPE Debug) + +find_package(CUDA REQUIRED) + +include_directories(${PROJECT_SOURCE_DIR}/include) +# include and link dirs of cuda and tensorrt, you need adapt them if yours are different +# cuda +include_directories(/usr/local/cuda-10.2/include) +link_directories(/usr/local/cuda-10.2/lib64) +# tensorrt +include_directories(/home/jushi/TensorRT-7.2.1.6/include) +link_directories(/home/jushi/TensorRT-7.2.1.6/lib) + +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED") + +find_package(OpenCV) +include_directories(${OpenCV_INCLUDE_DIRS}) + +add_executable(detr ${PROJECT_SOURCE_DIR}/detr.cpp) +target_link_libraries(detr nvinfer) +target_link_libraries(detr cudart) +target_link_libraries(detr ${OpenCV_LIBS}) + +add_definitions(-O2 -pthread) + diff --git a/detr/README.md b/detr/README.md new file mode 100644 index 0000000..0e6d963 --- /dev/null +++ b/detr/README.md @@ -0,0 +1,58 @@ +# DETR + +The Pytorch implementation is [facebookresearch/detr](https://github.com/facebookresearch/detr). + +For details see [End-to-End Object Detection with Transformers](https://ai.facebook.com/research/publications/end-to-end-object-detection-with-transformers). + +## Test Environment + +- GTX2080Ti / Ubuntu16.04 / cuda10.2 / cudnn8.0.4 / TensorRT7.2.1 / OpenCV4.2 +- GTX2080Ti / win10 / cuda10.2 / cudnn8.0.4 / TensorRT7.2.1 / OpenCV4.2 / VS2017 + +## How to Run + +1. generate .wts from pytorch with .pth + +``` +// git clone https://github.com/facebookresearch/detr.git +// go to facebookresearch/detr +// download https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth +// download https://raw.githubusercontent.com/freedenS/TestImage/main/demo.jpg +// copy tensorrtx/detr/gen_wts.py and demo.jpg into facebookresearch/detr +python gen_wts.py +// a file 'detr.wts' will be generated. +``` + +2. build tensorrtx/detr and run + +``` +// put detr.wts into tensorrtx/detr +// go to tensorrtx/detr +// update parameters in detr.cpp if your model is trained on custom dataset.The parameters are corresponding to config in detr. +mkdir build +cd build +cmake .. +make +sudo ./detr -s [.wts] // serialize model to plan file +sudo ./detr -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed +// For example +sudo ./detr -s ../detr.wts detr.engine +sudo ./detr -d detr.engine ../samples +``` + +3. check the images generated, as follows. _demo.jpg and so on. + +## NOTE + +- tensorrt use fixed input size, if the size of your data is different from the engine, you need to adjust your data and the result. +- image preprocessing with c++ is a little different with python(opencv vs PIL) + + +## Latency + +average cost of doInference(in detr.cpp) from second time with batch=1 under the ubuntu environment above + +| | fp32 | fp16 | int8 | +| ---- | ------- | ------- | ---- | +| R50 | 19.57ms | 9.424ms | TODO | + diff --git a/detr/backbone.hpp b/detr/backbone.hpp new file mode 100644 index 0000000..3b4c3d4 --- /dev/null +++ b/detr/backbone.hpp @@ -0,0 +1,326 @@ +#pragma once +#include +#include "common.hpp" + +enum RESNETTYPE { + R18 = 0, + R34, + R50, + R101, + R152 +}; + +const std::map> num_blocks_per_stage = { + {R18, {2, 2, 2, 2}}, + {R34, {3, 4, 6, 3}}, + {R50, {3, 4, 6, 3}}, + {R101, {3, 4, 23, 3}}, + {R152, {3, 8, 36, 3}} +}; + +IScaleLayer* addBatchNorm2d( +INetworkDefinition *network, +std::unordered_map& weightMap, +ITensor& input, +const std::string& lname, +float eps = 1e-5 +) { + float *gamma = (float*)(weightMap[lname + ".weight"].values); + float *beta = (float*)(weightMap[lname + ".bias"].values); + float *mean = (float*)(weightMap[lname + ".running_mean"].values); + float *var = (float*)(weightMap[lname + ".running_var"].values); + int len = weightMap[lname + ".running_var"].count; + + float *scval = reinterpret_cast(malloc(sizeof(float) * len)); + for (int i = 0; i < len; i++) { + scval[i] = gamma[i] / sqrt(var[i] + eps); + } + Weights scale{ DataType::kFLOAT, scval, len }; + + float *shval = reinterpret_cast(malloc(sizeof(float) * len)); + for (int i = 0; i < len; i++) { + shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps); + } + Weights shift{ DataType::kFLOAT, shval, len }; + + float *pval = reinterpret_cast(malloc(sizeof(float) * len)); + for (int i = 0; i < len; i++) { + pval[i] = 1.0; + } + Weights power{ DataType::kFLOAT, pval, len }; + + weightMap[lname + ".scale"] = scale; + weightMap[lname + ".shift"] = shift; + weightMap[lname + ".power"] = power; + IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power); + assert(scale_1); + return scale_1; +} + +ILayer* BasicStem( +INetworkDefinition *network, +std::unordered_map& weightMap, +const std::string& lname, +ITensor& input, +int out_channels, +int group_num = 1 +) { + // conv1 + Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; + IConvolutionLayer* conv1 = network->addConvolutionNd( + input, + out_channels, + DimsHW{ 7, 7 }, + weightMap[lname + ".conv1.weight"], + emptywts); + assert(conv1); + conv1->setStrideNd(DimsHW{ 2, 2 }); + conv1->setPaddingNd(DimsHW{ 3, 3 }); + conv1->setNbGroups(group_num); + + auto bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn1"); + assert(bn1); + + auto r1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU); + assert(r1); + + auto max_pool2d = network->addPoolingNd(*r1->getOutput(0), PoolingType::kMAX, DimsHW{ 3, 3 }); + max_pool2d->setStrideNd(DimsHW{ 2, 2 }); + max_pool2d->setPaddingNd(DimsHW{ 1, 1 }); + auto mp_dim = max_pool2d->getOutput(0)->getDimensions(); + return max_pool2d; +} + +ITensor* BasicBlock( +INetworkDefinition *network, +std::unordered_map& weightMap, +const std::string& lname, +ITensor& input, +int in_channels, +int out_channels, +int stride = 1 +) { + // conv1 + IConvolutionLayer* conv1 = network->addConvolutionNd( + input, + out_channels, + DimsHW{ 3, 3 }, + weightMap[lname + ".conv1.weight"], + weightMap[lname + ".conv1.bias"]); + assert(conv1); + conv1->setStrideNd(DimsHW{ stride, stride }); + conv1->setPaddingNd(DimsHW{ 1, 1 }); + + auto r1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); + assert(r1); + + // conv2 + IConvolutionLayer* conv2 = network->addConvolutionNd( + *r1->getOutput(0), + out_channels, DimsHW{ 3, 3 }, + weightMap[lname + ".conv2.weight"], + weightMap[lname + ".conv2.bias"]); + assert(conv2); + conv2->setStrideNd(DimsHW{ 1, 1 }); + conv2->setPaddingNd(DimsHW{ 1, 1 }); + + // shortcut + ITensor* shortcut_value = nullptr; + if (in_channels != out_channels) { + auto shortcut = network->addConvolutionNd( + input, + out_channels, + DimsHW{ 1, 1 }, + weightMap[lname + ".shortcut.weight"], + weightMap[lname + ".shortcut.bias"]); + assert(shortcut); + shortcut->setStrideNd(DimsHW{ stride, stride }); + shortcut_value = shortcut->getOutput(0); + } else { + shortcut_value = &input; + } + + // add + auto ew = network->addElementWise(*conv2->getOutput(0), *shortcut_value, ElementWiseOperation::kSUM); + assert(ew); + + auto r3 = network->addActivation(*ew->getOutput(0), ActivationType::kRELU); + assert(r3); + + return r3->getOutput(0); +} + +ITensor* BottleneckBlock( +INetworkDefinition *network, +std::unordered_map& weightMap, +const std::string& lname, +ITensor& input, +int in_channels, +int bottleneck_channels, +int out_channels, +int stride = 1, +int dilation = 1, +int group_num = 1 +) { + Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; + // conv1 + IConvolutionLayer* conv1 = network->addConvolutionNd( + input, + bottleneck_channels, + DimsHW{ 1, 1 }, + weightMap[lname + ".conv1.weight"], + emptywts); + assert(conv1); + conv1->setStrideNd(DimsHW{ 1, 1 }); + conv1->setNbGroups(group_num); + + auto bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn1"); + assert(bn1); + + auto r1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU); + assert(r1); + + // conv2 + IConvolutionLayer* conv2 = network->addConvolutionNd( + *r1->getOutput(0), + bottleneck_channels, + DimsHW{ 3, 3 }, + weightMap[lname + ".conv2.weight"], + emptywts); + assert(conv2); + conv2->setStrideNd(DimsHW{ stride, stride }); + conv2->setPaddingNd(DimsHW{ 1 * dilation, 1 * dilation }); + conv2->setDilationNd(DimsHW{ dilation, dilation }); + conv2->setNbGroups(group_num); + + auto bn2 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + ".bn2"); + assert(bn2); + + auto r2 = network->addActivation(*bn2->getOutput(0), ActivationType::kRELU); + assert(r2); + + // conv3 + IConvolutionLayer* conv3 = network->addConvolutionNd( + *r2->getOutput(0), + out_channels, + DimsHW{ 1, 1 }, + weightMap[lname + ".conv3.weight"], + emptywts); + assert(conv3); + conv3->setStrideNd(DimsHW{ 1, 1 }); + conv3->setNbGroups(group_num); + + auto bn3 = addBatchNorm2d(network, weightMap, *conv3->getOutput(0), lname + ".bn3"); + assert(bn3); + + // shortcut + ITensor* shortcut_value = nullptr; + if (in_channels != out_channels) { + auto shortcut = network->addConvolutionNd( + input, + out_channels, + DimsHW{ 1, 1 }, + weightMap[lname + ".downsample.0.weight"], + emptywts); + assert(shortcut); + shortcut->setStrideNd(DimsHW{stride, stride}); + shortcut->setNbGroups(group_num); + + auto shortcut_bn = addBatchNorm2d(network, weightMap, *shortcut->getOutput(0), lname + ".downsample.1"); + assert(shortcut_bn); + shortcut_value = shortcut_bn->getOutput(0); + } else { + shortcut_value = &input; + } + + // add + auto ew = network->addElementWise(*bn3->getOutput(0), *shortcut_value, ElementWiseOperation::kSUM); + assert(ew); + + auto r3 = network->addActivation(*ew->getOutput(0), ActivationType::kRELU); + assert(r3); + + return r3->getOutput(0); +} + +ITensor* MakeStage( +INetworkDefinition *network, +std::unordered_map& weightMap, +const std::string& lname, +ITensor& input, +int stage, +RESNETTYPE resnet_type, +int in_channels, +int bottleneck_channels, +int out_channels, +int first_stride = 1, +int dilation = 1 +) { + ITensor* out = &input; + for (int i = 0; i < stage; i++) { + std::string layerName = lname + "." + std::to_string(i); + int stride = i == 0 ? first_stride : 1; + + if (resnet_type == R18 || resnet_type == R34) + out = BasicBlock(network, weightMap, layerName, *out, in_channels, out_channels, stride); + else + out = BottleneckBlock( + network, + weightMap, + layerName, + *out, + in_channels, + bottleneck_channels, + out_channels, + stride, + dilation); + + in_channels = out_channels; + } + return out; +} + +ITensor* BuildResNet( +INetworkDefinition *network, +std::unordered_map& weightMap, +ITensor& input, +RESNETTYPE resnet_type, +int stem_out_channels, +int bottleneck_channels, +int res2_out_channels, +int res5_dilation = 1 +) { + assert(res5_dilation == 1 || res5_dilation == 2); // "res5_dilation must be 1 or 2" + if (resnet_type == R18 || resnet_type == R34) { + assert(res2_out_channels == 64); // "res2_out_channels must be 64 for R18/R34") + assert(res5_dilation == 1); // "res5_dilation must be 1 for R18/R34") + } + + int out_channels = res2_out_channels; + ITensor* out = nullptr; + // stem + auto stem = BasicStem(network, weightMap, "backbone.0.body", input, stem_out_channels); + out = stem->getOutput(0); + + // res + for (int i = 0; i < 4; i++) { + int dilation = (i == 3) ? res5_dilation : 1; + int first_stride = (i == 0 || (i == 3 && dilation == 2)) ? 1 : 2; + out = MakeStage( + network, + weightMap, + "backbone.0.body.layer" + std::to_string(i + 1), + *out, + num_blocks_per_stage.at(resnet_type)[i], + resnet_type, + stem_out_channels, + bottleneck_channels, + out_channels, + first_stride, + dilation); + stem_out_channels = out_channels; + bottleneck_channels *= 2; + out_channels *= 2; + } + return out; +} diff --git a/detr/common.hpp b/detr/common.hpp new file mode 100644 index 0000000..7a3550d --- /dev/null +++ b/detr/common.hpp @@ -0,0 +1,90 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include "./logging.h" +#include +#include + +static Logger gLogger; + +using namespace nvinfer1; +void loadWeights(const std::string file, std::unordered_map& weightMap) { + std::cout << "Loading weights: " << file << std::endl; + + // Open weights file + std::ifstream input(file); + assert(input.is_open() && "Unable to load weight file. please check if the .wts file path is right!!!!!!"); + + // Read number of weight blobs + int32_t count; + input >> count; + assert(count > 0 && "Invalid weight map file."); + + while (count--) { + Weights wt{ DataType::kFLOAT, nullptr, 0 }; + uint32_t size; + + // Read name and type of blob + std::string name; + input >> name >> std::dec >> size; + wt.type = DataType::kFLOAT; + + // Load blob + uint32_t* val = reinterpret_cast(malloc(sizeof(val) * size)); + for (uint32_t x = 0, y = size; x < y; ++x) { + input >> std::hex >> val[x]; + } + wt.values = val; + + wt.count = size; + weightMap[name] = wt; + } +} + +int CalculateSize(Dims a) { + int res = 1; + for (int i = 0; i < a.nbDims; i++) { + res *= a.d[i]; + } + return res; +} + +static inline int read_files_in_dir(const char *p_dir_name, std::vector &file_names) { + DIR *p_dir = opendir(p_dir_name); + if (p_dir == nullptr) { + return -1; + } + + struct dirent* p_file = nullptr; + while ((p_file = readdir(p_dir)) != nullptr) { + if (strcmp(p_file->d_name, ".") != 0 && + strcmp(p_file->d_name, "..") != 0) { + // std::string cur_file_name(p_dir_name); + // cur_file_name += "/"; + // cur_file_name += p_file->d_name; + std::string cur_file_name(p_file->d_name); + file_names.push_back(cur_file_name); + } + } + + closedir(p_dir); + return 0; +} + +#ifndef CUDA_CHECK +#define CUDA_CHECK(callstr)\ + {\ + cudaError_t error_code = callstr;\ + if (error_code != cudaSuccess) {\ + std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__;\ + assert(0);\ + }\ + } +#endif // CUDA_CHECK diff --git a/detr/detr.cpp b/detr/detr.cpp new file mode 100644 index 0000000..00c537d --- /dev/null +++ b/detr/detr.cpp @@ -0,0 +1,826 @@ +#pragma once +#include +#include +#include "./logging.h" +#include "backbone.hpp" + +#define DEVICE 0 +#define BATCH_SIZE 1 + +// 1 / math.sqrt(head_dim) https://github.com/pytorch/pytorch/blob/master/torch/csrc/api/include/torch/nn/functional/activation.h#623 +static const float SCALING = 0.17677669529663687; +static const float MIN_SIZE = 800.0; +static const int INPUT_H = 800; +static const int INPUT_W = 1066; +static const int NUM_CLASS = 92; // include background +static const float SCALING_ONE = 1.0; +static const float SHIFT_ZERO = 0.0; +static const float POWER_TWO = 2.0; +static const float EPS = 0.00001; +static const int D_MODEL = 256; +static const int NHEAD = 8; +static const int DIM_FEEDFORWARD = 2048; +static const int NUM_ENCODE_LAYERS = 6; +static const int NUM_DECODE_LAYERS = 6; +static const int NUM_QUERIES = 100; +static const float SCORE_THRESH = 0.5; + +const char* INPUT_NODE_NAME = "images"; +const std::vector OUTPUT_NAMES = { "scores", "boxes"}; + +void preprocessImg(cv::Mat& img) { + // convert to rgb + cv::cvtColor(img, img, cv::COLOR_BGR2RGB); + float ratio = static_cast(MIN_SIZE) / std::min(img.rows, img.cols); + int newh = 0, neww = 0; + if (img.rows < img.cols) { + newh = MIN_SIZE; + neww = ratio * img.cols; + } else { + newh = ratio * img.rows; + neww = MIN_SIZE; + } + cv::resize(img, img, cv::Size(neww, newh)); + img.convertTo(img, CV_32FC3); + img /= 255; + img -= cv::Scalar(0.485, 0.456, 0.406); + img /= cv::Scalar(0.229, 0.224, 0.225); +} + +ITensor* PositionEmbeddingSine( +INetworkDefinition *network, +std::unordered_map& weightMap, +ITensor& input, +int num_pos_feats = 64, +int temperature = 10000 +) { + // refer to https://github.com/facebookresearch/detr/blob/master/models/position_encoding.py#12 + // TODO: improve this implementation + auto mask_dim = input.getDimensions(); + int h = mask_dim.d[1], w = mask_dim.d[2]; + std::vector> y_embed(h); + for (int i = 0; i < h; i++) + y_embed[i] = std::vector(w, i + 1); + std::vector sub_embed(w, 0); + for (int i = 0; i < w; i++) + sub_embed[i] = i + 1; + std::vector> x_embed(h, sub_embed); + + // normalize + float eps = 1e-6, scale = 2.0 * 3.1415926; + for (int i = 0; i < h; i++) { + for (int j = 0; j < w; j++) { + y_embed[i][j] = y_embed[i][j] / (h + eps) * scale; + x_embed[i][j] = x_embed[i][j] / (w + eps) * scale; + } + } + + // dim_t + std::vector dim_t(num_pos_feats, 0); + for (int i = 0; i < num_pos_feats; i++) { + dim_t[i] = pow(temperature, (2 * (i / 2) / static_cast(num_pos_feats))); + } + + // pos_x, pos_y + std::vector>> pos_x(h, + std::vector>(w, + std::vector(num_pos_feats, 0))); + + std::vector>> pos_y(h, + std::vector>(w, + std::vector(num_pos_feats, 0))); + for (int i = 0; i < h; i++) { + for (int j = 0; j < w; j++) { + for (int k = 0; k < num_pos_feats; k++) { + float value_x = x_embed[i][j] / dim_t[k]; + float value_y = y_embed[i][j] / dim_t[k]; + if (k & 1) { + pos_x[i][j][k] = std::cos(value_x); + pos_y[i][j][k] = std::cos(value_y); + } else { + pos_x[i][j][k] = std::sin(value_x); + pos_y[i][j][k] = std::sin(value_y); + } + } + } + } + + // pos + float *pval = reinterpret_cast(malloc(sizeof(float) * h * w * num_pos_feats * 2)); + float *pNext = pval; + for (int i = 0; i < h; i++) { + for (int j = 0; j < w; j++) { + for (int k = 0; k < num_pos_feats; k++) { + *pNext = pos_y[i][j][k]; + ++pNext; + } + for (int k = 0; k < num_pos_feats; k++) { + *pNext = pos_x[i][j][k]; + ++pNext; + } + } + } + Weights pos_embed_weight{ DataType::kFLOAT, pval, h * w * num_pos_feats * 2 }; + weightMap["pos"] = pos_embed_weight; + auto pos_embed = network->addConstant(Dims4{ h * w, num_pos_feats * 2, 1, 1 }, pos_embed_weight); + assert(pos_embed); + return pos_embed->getOutput(0); +} + +ITensor* MultiHeadAttention( +INetworkDefinition *network, +std::unordered_map& weightMap, +const std::string& lname, +ITensor& query, +ITensor& key, +ITensor& value, +int embed_dim = 256, +int num_heads = 8 +) { + int tgt_len = query.getDimensions().d[0]; + int head_dim = embed_dim / num_heads; + + // q + auto linear_q = network->addFullyConnected( + query, + embed_dim, + weightMap[lname + ".in_proj_weight_q"], + weightMap[lname + ".in_proj_bias_q"]); + assert(linear_q); + + // k + auto linear_k = network->addFullyConnected( + key, + embed_dim, + weightMap[lname + ".in_proj_weight_k"], + weightMap[lname + ".in_proj_bias_k"]); + assert(linear_k); + + // v + auto linear_v = network->addFullyConnected( + value, + embed_dim, + weightMap[lname + ".in_proj_weight_v"], + weightMap[lname + ".in_proj_bias_v"]); + assert(linear_v); + + auto scaling_t = network->addConstant(Dims4{ 1, 1, 1, 1 }, Weights{ DataType::kFLOAT, &SCALING, 1 }); + assert(scaling_t); + auto q_scaling = network->addElementWise( + *linear_q->getOutput(0), + *scaling_t->getOutput(0), + ElementWiseOperation::kPROD); + assert(q_scaling); + + auto q_shuffle = network->addShuffle(*q_scaling->getOutput(0)); + assert(q_shuffle); + q_shuffle->setName((lname + ".q_shuffle").c_str()); + q_shuffle->setReshapeDimensions(Dims3{ -1, num_heads, head_dim }); + q_shuffle->setSecondTranspose(Permutation{1, 0, 2}); + + auto k_shuffle = network->addShuffle(*linear_k->getOutput(0)); + assert(k_shuffle); + k_shuffle->setName((lname + ".k_shuffle").c_str()); + k_shuffle->setReshapeDimensions(Dims3{ -1, num_heads, head_dim }); + k_shuffle->setSecondTranspose(Permutation{ 1, 0, 2 }); + + auto v_shuffle = network->addShuffle(*linear_v->getOutput(0)); + assert(v_shuffle); + v_shuffle->setName((lname + ".v_shuffle").c_str()); + v_shuffle->setReshapeDimensions(Dims3{ -1, num_heads, head_dim }); + v_shuffle->setSecondTranspose(Permutation{ 1, 0, 2 }); + + auto q_product_k = network->addMatrixMultiply(*q_shuffle->getOutput(0), false, *k_shuffle->getOutput(0), true); + assert(q_product_k); + + // src_key_padding_mask are all false, so do nothing here + // see https://github.com/pytorch/pytorch/blob/master/torch/csrc/api/include/torch/nn/functional/activation.h#826-#839 + + auto softmax = network->addSoftMax(*q_product_k->getOutput(0)); + assert(softmax); + softmax->setAxes(4); + + auto attn_product_v = network->addMatrixMultiply(*softmax->getOutput(0), false, *v_shuffle->getOutput(0), false); + assert(attn_product_v); + + auto attn_shuffle = network->addShuffle(*attn_product_v->getOutput(0)); + assert(attn_shuffle); + attn_shuffle->setName((lname + ".attn_shuffle").c_str()); + attn_shuffle->setFirstTranspose(Permutation{ 1, 0, 2 }); + attn_shuffle->setReshapeDimensions(Dims4{ tgt_len, -1, 1, 1 }); + + auto linear_attn = network->addFullyConnected( + *attn_shuffle->getOutput(0), + embed_dim, + weightMap[lname + ".out_proj.weight"], + weightMap[lname + ".out_proj.bias"]); + assert(linear_attn); + + return linear_attn->getOutput(0); +} + +ITensor* LayerNorm( +INetworkDefinition *network, +ITensor& input, +std::unordered_map& weightMap, +const std::string& lname, +int d_model = 256 +) { + // TODO: maybe a better implementation https://github.com/NVIDIA/TensorRT/blob/master/plugin/common/common.cuh#212 + auto mean = network->addReduce(input, ReduceOperation::kAVG, 2, true); + assert(mean); + + auto sub_mean = network->addElementWise(input, *mean->getOutput(0), ElementWiseOperation::kSUB); + assert(sub_mean); + + // implement pow2 with scale + Weights scale{ DataType::kFLOAT, &SCALING_ONE, 1 }; + Weights shift{ DataType::kFLOAT, &SHIFT_ZERO, 1 }; + Weights power{ DataType::kFLOAT, &POWER_TWO, 1 }; + auto pow2 = network->addScaleNd(*sub_mean->getOutput(0), ScaleMode::kUNIFORM, shift, scale, power, 0); + assert(pow2); + + auto pow_mean = network->addReduce(*pow2->getOutput(0), ReduceOperation::kAVG, 2, true); + assert(pow_mean); + + auto eps = network->addConstant(Dims4{ 1, 1, 1, 1 }, Weights{ DataType::kFLOAT, &EPS, 1 }); + assert(eps); + + auto add_eps = network->addElementWise(*pow_mean->getOutput(0), *eps->getOutput(0), ElementWiseOperation::kSUM); + assert(add_eps); + + auto sqrt = network->addUnary(*add_eps->getOutput(0), UnaryOperation::kSQRT); + assert(sqrt); + + auto div = network->addElementWise(*sub_mean->getOutput(0), *sqrt->getOutput(0), ElementWiseOperation::kDIV); + assert(div); + + float *pval = reinterpret_cast(malloc(sizeof(float) * d_model)); + for (int i = 0; i < d_model; i++) { + pval[i] = 1.0; + } + Weights norm1_power{ DataType::kFLOAT, pval, d_model }; + weightMap[lname + ".power"] = norm1_power; + auto affine = network->addScaleNd( + *div->getOutput(0), + ScaleMode::kCHANNEL, + weightMap[lname + ".bias"], + weightMap[lname + ".weight"], + norm1_power, + 1); + assert(affine); + return affine->getOutput(0); +} + +ITensor* TransformerEncoderLayer( +INetworkDefinition *network, +std::unordered_map& weightMap, +const std::string& lname, +ITensor& src, +ITensor& pos, +int d_model = 256, +int nhead = 8, +int dim_feedforward = 2048 +) { + auto pos_embed = network->addElementWise(src, pos, ElementWiseOperation::kSUM); + assert(pos_embed); + + ITensor* src2 = MultiHeadAttention( + network, + weightMap, + lname + ".self_attn", + *pos_embed->getOutput(0), + *pos_embed->getOutput(0), + src, + d_model, + nhead); + + auto shortcut1 = network->addElementWise(src, *src2, ElementWiseOperation::kSUM); + assert(shortcut1); + + ITensor* norm1 = LayerNorm(network, *shortcut1->getOutput(0), weightMap, lname + ".norm1"); + + auto linear1 = network->addFullyConnected( + *norm1, + dim_feedforward, + weightMap[lname + ".linear1.weight"], + weightMap[lname + ".linear1.bias"]); + assert(linear1); + + auto relu = network->addActivation(*linear1->getOutput(0), ActivationType::kRELU); + assert(relu); + + auto linear2 = network->addFullyConnected( + *relu->getOutput(0), + d_model, + weightMap[lname + ".linear2.weight"], + weightMap[lname + ".linear2.bias"]); + assert(linear2); + + auto shortcut2 = network->addElementWise(*norm1, *linear2->getOutput(0), ElementWiseOperation::kSUM); + assert(shortcut2); + + ITensor* norm2 = LayerNorm(network, *shortcut2->getOutput(0), weightMap, lname + ".norm2"); + return norm2; +} + +ITensor* TransformerEncoder( +INetworkDefinition *network, +std::unordered_map& weightMap, +const std::string& lname, +ITensor& src, +ITensor& pos, +int num_layers = 6 +) { + ITensor* out = &src; + for (int i = 0; i < num_layers; i++) { + std::string layer_name = lname + ".layers." + std::to_string(i); + out = TransformerEncoderLayer(network, weightMap, layer_name, *out, pos); + } + return out; +} + +ITensor* TransformerDecoderLayer( +INetworkDefinition *network, +std::unordered_map& weightMap, +const std::string& lname, +ITensor& tgt, +ITensor& memory, +ITensor& pos, +ITensor& query_pos, +int d_model = 256, +int nhead = 8, +int dim_feedforward = 2048 +) { + auto pos_embed = network->addElementWise(tgt, query_pos, ElementWiseOperation::kSUM); + assert(pos_embed); + + ITensor* tgt2 = MultiHeadAttention( + network, + weightMap, + lname + ".self_attn", + *pos_embed->getOutput(0), + *pos_embed->getOutput(0), + tgt); + + auto shortcut1 = network->addElementWise(tgt, *tgt2, ElementWiseOperation::kSUM); + assert(shortcut1); + ITensor* norm1 = LayerNorm(network, *shortcut1->getOutput(0), weightMap, lname + ".norm1"); + + auto query_embed = network->addElementWise(*norm1, query_pos, ElementWiseOperation::kSUM); + assert(query_embed); + + auto key_embed = network->addElementWise(memory, pos, ElementWiseOperation::kSUM); + assert(key_embed); + + ITensor* mha2 = MultiHeadAttention( + network, + weightMap, + lname + ".multihead_attn", + *query_embed->getOutput(0), + *key_embed->getOutput(0), + memory); + + auto shortcut2 = network->addElementWise(*norm1, *mha2, ElementWiseOperation::kSUM); + assert(shortcut2); + + ITensor* norm2 = LayerNorm(network, *shortcut2->getOutput(0), weightMap, lname + ".norm2"); + + auto linear1 = network->addFullyConnected( + *norm2, + dim_feedforward, + weightMap[lname + ".linear1.weight"], + weightMap[lname + ".linear1.bias"]); + assert(linear1); + + auto relu = network->addActivation(*linear1->getOutput(0), ActivationType::kRELU); + assert(relu); + + auto linear2 = network->addFullyConnected( + *relu->getOutput(0), + d_model, + weightMap[lname + ".linear2.weight"], + weightMap[lname + ".linear2.bias"]); + assert(linear2); + + auto shortcut3 = network->addElementWise(*norm2, *linear2->getOutput(0), ElementWiseOperation::kSUM); + assert(shortcut3); + + ITensor* norm3 = LayerNorm(network, *shortcut3->getOutput(0), weightMap, lname + ".norm3"); + + return norm3; +} + +ITensor* TransformerDecoder( +INetworkDefinition *network, +std::unordered_map& weightMap, +const std::string& lname, +ITensor& tgt, +ITensor& memory, +ITensor& pos, +ITensor& query_pos, +int num_layers = 6, +int d_model = 256, +int nhead = 8, +int dim_feedforward = 2048 +) { + ITensor* out = &tgt; + for (int i = 0; i < num_layers; i++) { + std::string layer_name = lname + ".layers." + std::to_string(i); + out = TransformerDecoderLayer( + network, + weightMap, + layer_name, + *out, + memory, + pos, + query_pos, + d_model, + nhead, + dim_feedforward); + } + ITensor* norm = LayerNorm(network, *out, weightMap, lname + ".norm", d_model); + return norm; +} + +ITensor* Transformer( +INetworkDefinition *network, +std::unordered_map& weightMap, +const std::string& lname, +ITensor& src, +ITensor& pos_embed, +int num_queries = 100, +int num_encoder_layers = 6, +int num_decoder_layers = 6, +int d_model = 256, +int nhead = 8, +int dim_feedforward = 2048 +) { + auto memory = TransformerEncoder(network, weightMap, lname + ".encoder", src, pos_embed, num_encoder_layers); + + // construct tgt + float *pval = reinterpret_cast(malloc(sizeof(float) * num_queries * d_model)); + for (int i = 0; i < num_queries * d_model; i++) { + pval[i] = 0.0; + } + Weights tgt_weight{ DataType::kFLOAT, pval, num_queries * d_model }; + weightMap[lname + ".tgt_weight"] = tgt_weight; + auto tgt = network->addConstant(Dims4{ num_queries, d_model, 1, 1 }, tgt_weight); + assert(tgt); + // construct query_pos + auto query_pos = network->addConstant(Dims4{ num_queries, d_model, 1, 1 }, weightMap["query_embed.weight"]); + assert(query_pos); + + auto out = TransformerDecoder( + network, + weightMap, + lname + ".decoder", + *tgt->getOutput(0), + *memory, pos_embed, + *query_pos->getOutput(0), + num_decoder_layers, + d_model, + nhead, + dim_feedforward); + return out; +} + +ITensor* MLP( +INetworkDefinition *network, +std::unordered_map& weightMap, +const std::string& lname, +ITensor& src, +int num_layers = 3, +int hidden_dim = 256, +int output_dim = 4 +) { + ITensor* out = &src; + for (int i = 0; i < num_layers; i++) { + std::string layer_name = lname + "." + std::to_string(i); + if (i != num_layers - 1) { + auto fc = network->addFullyConnected( + *out, + hidden_dim, + weightMap[layer_name + ".weight"], + weightMap[layer_name + ".bias"]); + assert(fc); + auto relu = network->addActivation(*fc->getOutput(0), ActivationType::kRELU); + assert(relu); + out = relu->getOutput(0); + } else { + auto fc = network->addFullyConnected( + *out, + output_dim, + weightMap[layer_name + ".weight"], + weightMap[layer_name + ".bias"]); + assert(fc); + out = fc->getOutput(0); + } + } + return out; +} + +std::vector Predict( +INetworkDefinition *network, +std::unordered_map& weightMap, +ITensor& src +) { + auto class_embed = network->addFullyConnected( + src, + NUM_CLASS, + weightMap["class_embed.weight"], + weightMap["class_embed.bias"]); + assert(class_embed); + auto class_softmax = network->addSoftMax(*class_embed->getOutput(0)); + assert(class_softmax); + class_softmax->setAxes(2); + ITensor* bbox = MLP(network, weightMap, "bbox_embed.layers", src); + auto bbox_sig = network->addActivation(*bbox, ActivationType::kSIGMOID); + assert(bbox_sig); + std::vector output = { class_softmax->getOutput(0), bbox_sig->getOutput(0) }; + return output; +} + +ICudaEngine* createEngine_r50detr( +unsigned int maxBatchSize, +const std::string& wtsfile, +IBuilder* builder, +IBuilderConfig* config, +DataType dt, +const std::string& modelType = "fp16" +) { + /* + description: after fuse bn + */ + INetworkDefinition* network = builder->createNetworkV2(0U); + + // Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME + ITensor* data = network->addInput("data", dt, Dims3{ 3, INPUT_H, INPUT_W }); + + // preprocess + std::unordered_map weightMap; + loadWeights(wtsfile, weightMap); + + // backbone + auto features = BuildResNet(network, weightMap, *data, R50, 64, 64, 256); + ITensor* pos_embed = PositionEmbeddingSine(network, weightMap, *features, 128); + auto input_proj = network->addConvolutionNd( + *features, + D_MODEL, + DimsHW{ 1, 1 }, + weightMap["input_proj.weight"], + weightMap["input_proj.bias"]); + assert(input_proj); + input_proj->setStrideNd(DimsHW{ 1, 1 }); + auto flatten = network->addShuffle(*input_proj->getOutput(0)); + assert(flatten); + flatten->setReshapeDimensions(Dims4{ input_proj->getOutput(0)->getDimensions().d[0], -1, 1, 1 }); + flatten->setSecondTranspose(Permutation{ 1, 0, 2, 3 }); + + auto out1 = Transformer( + network, + weightMap, + "transformer", + *flatten->getOutput(0), + *pos_embed, + NUM_QUERIES, + NUM_ENCODE_LAYERS, + NUM_DECODE_LAYERS, + D_MODEL, + NHEAD, + DIM_FEEDFORWARD); + std::vector results = Predict(network, weightMap, *out1); + + // build output + for (int i = 0; i < results.size(); i++) { + network->markOutput(*results[i]); + results[i]->setName(OUTPUT_NAMES[i].c_str()); + } + + // build engine + builder->setMaxBatchSize(maxBatchSize); + config->setMaxWorkspaceSize(1ULL << 30); + + if (modelType == "fp32") { + } else if (modelType == "fp16") { + config->setFlag(BuilderFlag::kFP16); + } else if (modelType == "int8") { + // TODO: test with int8 quantization + } else { + throw("does not support model type"); + } + + std::cout << "Building engine, please wait for a while..." << std::endl; + ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); + std::cout << "Build engine successfully!" << std::endl; + + // destroy network + network->destroy(); + + // Release host memory + for (auto& mem : weightMap) { + free((void*)(mem.second.values)); + } + return engine; +} + +void BuildDETRModel(unsigned int maxBatchSize, IHostMemory** modelStream, +const std::string& wtsfile, std::string modelType = "fp32") { + // Create builder + IBuilder* builder = createInferBuilder(gLogger); + IBuilderConfig* config = builder->createBuilderConfig(); + + // Create model to populate the network, then set the outputs and create an engine + ICudaEngine* engine = createEngine_r50detr(maxBatchSize, + wtsfile, builder, config, DataType::kFLOAT, modelType); + assert(engine != nullptr); + + // Serialize the engine + (*modelStream) = engine->serialize(); + + // Close everything down + engine->destroy(); + builder->destroy(); +} + +void doInference(IExecutionContext& context, cudaStream_t& stream, std::vector& buffers, +std::vector& input, std::vector& output) { + CUDA_CHECK(cudaMemcpyAsync(buffers[0], input.data(), input.size() * sizeof(float), + cudaMemcpyHostToDevice, stream)); + + context.enqueue(BATCH_SIZE, buffers.data(), stream, nullptr); + + CUDA_CHECK(cudaMemcpyAsync(output[0], buffers[1], BATCH_SIZE * NUM_QUERIES * NUM_CLASS * sizeof(float), + cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaMemcpyAsync(output[1], buffers[2], BATCH_SIZE * NUM_QUERIES * 4 * sizeof(float), + cudaMemcpyDeviceToHost, stream)); + + cudaStreamSynchronize(stream); +} + +bool parse_args(int argc, char** argv, std::string& wtsFile, std::string& engineFile, std::string& imgDir) { + if (argc < 4) return false; + if (std::string(argv[1]) == "-s") { + wtsFile = std::string(argv[2]); + engineFile = std::string(argv[3]); + } else if (std::string(argv[1]) == "-d") { + engineFile = std::string(argv[2]); + imgDir = std::string(argv[3]); + } else { + return false; + } + return true; +} + +int main(int argc, char** argv) { + cudaSetDevice(DEVICE); + + std::string wtsFile = ""; + std::string engineFile = ""; + + std::string imgDir; + if (!parse_args(argc, argv, wtsFile, engineFile, imgDir)) { + std::cerr << "arguments not right!" << std::endl; + std::cerr << "./detr -s [.wts] [.engine] // serialize model to plan file" << std::endl; + std::cerr << "./detr -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; + return -1; + } + + if (!wtsFile.empty()) { + IHostMemory* modelStream{ nullptr }; + BuildDETRModel(BATCH_SIZE, &modelStream, wtsFile, "fp32"); + assert(modelStream != nullptr); + std::ofstream p(engineFile, std::ios::binary); + if (!p) { + std::cerr << "could not open plan output file" << std::endl; + return -1; + } + p.write(reinterpret_cast(modelStream->data()), modelStream->size()); + modelStream->destroy(); + return 0; + } + + // deserialize the .engine and run inference + std::ifstream file(engineFile, std::ios::binary); + if (!file.good()) { + std::cerr << "read " << engineFile << " error!" << std::endl; + return -1; + } + + std::string trtModelStream; + size_t modelSize{ 0 }; + file.seekg(0, file.end); + modelSize = file.tellg(); + file.seekg(0, file.beg); + trtModelStream.resize(modelSize); + assert(!trtModelStream.empty()); + file.read(const_cast(trtModelStream.c_str()), modelSize); + file.close(); + + // build engine + std::cout << "build engine" << std::endl; + IRuntime* runtime = createInferRuntime(gLogger); + assert(runtime != nullptr); + ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream.c_str(), modelSize); + assert(engine != nullptr); + IExecutionContext* context = engine->createExecutionContext(); + assert(context != nullptr); + runtime->destroy(); + + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreate(&stream)); + + // prepare input file + std::vector fileList; + if (read_files_in_dir(imgDir.c_str(), fileList) < 0) { + std::cerr << "read_files_in_dir failed." << std::endl; + return -1; + } + + // calculate input size + int input_size = CalculateSize(context->getBindingDimensions(0)); + + // prepare input data + std::vector data(BATCH_SIZE * input_size, 0); + void *data_d, *scores_d, *boxes_d; + CUDA_CHECK(cudaMalloc(&data_d, BATCH_SIZE * input_size * sizeof(float))); + CUDA_CHECK(cudaMalloc(&scores_d, BATCH_SIZE * NUM_QUERIES * NUM_CLASS * sizeof(float))); + CUDA_CHECK(cudaMalloc(&boxes_d, BATCH_SIZE * NUM_QUERIES * 4 * sizeof(float))); + + std::vector scores_h(BATCH_SIZE * NUM_QUERIES * NUM_CLASS); + std::vector boxes_h(BATCH_SIZE * NUM_QUERIES * 4); + + std::vector buffers = { data_d, scores_d, boxes_d }; + std::vector outputs = {scores_h.data(), boxes_h.data()}; + + int fcount = 0; + int fileLen = fileList.size(); + for (int f = 0; f < fileLen; f++) { + fcount++; + if (fcount < BATCH_SIZE && f + 1 != fileLen) continue; + + for (int b = 0; b < fcount; b++) { + cv::Mat img = cv::imread(imgDir + "/" + fileList[f - fcount + 1 + b]); + preprocessImg(img); + assert(img.cols * img.rows * 3 == input_size); + if (img.empty()) continue; + for (int c = 0; c < 3; c++) { + for (int h = 0; h < img.rows; h++) { + for (int w = 0; w < img.cols; w++) { + data[b * input_size + + c * img.rows * img.cols + h * img.cols + w] = img.at(h, w)[c]; + } + } + } + } + + // Run inference + auto start = std::chrono::system_clock::now(); + + doInference(*context, stream, buffers, data, outputs); + + auto end = std::chrono::system_clock::now(); + std::cout << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; + + for (int b = 0; b < fcount; b++) { + cv::Mat img = cv::imread(imgDir + "/" + fileList[f - fcount + 1 + b]); + for (int i = 0; i < scores_h.size(); i += NUM_CLASS) { + int label = -1; + float score = -1; + for (int j = i; j < i + NUM_CLASS; j++) { + if (score < scores_h[j]) { + label = j; + score = scores_h[j]; + } + } + if (score > SCORE_THRESH && (label % NUM_CLASS != NUM_CLASS - 1)) { + int ind = label / NUM_CLASS; + label = label % NUM_CLASS; + float cx = boxes_h[ind * 4]; + float cy = boxes_h[ind * 4 + 1]; + float w = boxes_h[ind * 4 + 2]; + float h = boxes_h[ind * 4 + 3]; + float x1 = (cx - w / 2.0) * img.cols; + float y1 = (cy - h / 2.0) * img.rows; + float x2 = (cx + w / 2.0) * img.cols; + float y2 = (cy + h / 2.0) * img.rows; + cv::Rect r(x1, y1, x2 - x1, y2 - y1); + cv::rectangle(img, r, cv::Scalar(0x27, 0xC1, 0x36), 2); + cv::putText(img, std::to_string(label), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, + cv::Scalar(0xFF, 0xFF, 0xFF), 2); + } + } + cv::imwrite("_" + fileList[f - fcount + 1 + b], img); + } + fcount = 0; + } + + cudaStreamDestroy(stream); + CUDA_CHECK(cudaFree(data_d)); + CUDA_CHECK(cudaFree(scores_d)); + CUDA_CHECK(cudaFree(boxes_d)); + context->destroy(); + engine->destroy(); + + return 0; +} diff --git a/detr/gen_wts.py b/detr/gen_wts.py new file mode 100644 index 0000000..518e66b --- /dev/null +++ b/detr/gen_wts.py @@ -0,0 +1,123 @@ +import cv2 + +import torch +from models.transformer import Transformer +from models.position_encoding import PositionEmbeddingSine +from models.backbone import Backbone, Joiner +from models.detr import DETR +import torchvision.transforms as T +from PIL import Image +import struct + +def box_cxcywh_to_xyxy(x): + x_c, y_c, w, h = x.unbind(-1) + b = [(x_c - 0.5 * w), (y_c - 0.5 * h), + (x_c + 0.5 * w), (y_c + 0.5 * h)] + return torch.stack(b, dim=-1) + +def build_backbone(): + N_steps = 256 // 2 + position_embedding = PositionEmbeddingSine(N_steps, normalize=True) + train_backbone = True + return_interm_layers = False + backbone = Backbone('resnet50', train_backbone, return_interm_layers, False) + model = Joiner(backbone, position_embedding) + model.num_channels = backbone.num_channels + return model + +def gen_wts(model, filename): + f = open(filename + '.wts', 'w') + f.write('{}\n'.format(len(model.state_dict().keys()) + 72)) + for k, v in model.state_dict().items(): + if 'in_proj' in k: + dim = int(v.size(0) / 3) + q_weight = v[:dim].reshape(-1).cpu().numpy() + k_weight = v[dim:2*dim].reshape(-1).cpu().numpy() + v_weight = v[2*dim:].reshape(-1).cpu().numpy() + f.write('{} {} '.format(k + '_q', len(q_weight))) + for vv in q_weight: + f.write(' ') + f.write(struct.pack('>f', float(vv)).hex()) + f.write('\n') + + f.write('{} {} '.format(k + '_k', len(k_weight))) + for vv in k_weight: + f.write(' ') + f.write(struct.pack('>f', float(vv)).hex()) + f.write('\n') + + f.write('{} {} '.format(k + '_v', len(v_weight))) + for vv in v_weight: + f.write(' ') + f.write(struct.pack('>f', float(vv)).hex()) + f.write('\n') + else: + vr = v.reshape(-1).cpu().numpy() + f.write('{} {} '.format(k, len(vr))) + for vv in vr: + f.write(' ') + f.write(struct.pack('>f',float(vv)).hex()) + f.write('\n') + f.close() + +def main(): + num_classes = 91 + device = torch.device('cuda') + + backbone = build_backbone() + + transformer = Transformer( + d_model=256, + dropout=0.1, + nhead=8, + dim_feedforward=2048, + num_encoder_layers=6, + num_decoder_layers=6, + normalize_before=False, + return_intermediate_dec=True, + ) + + model = DETR( + backbone, + transformer, + num_classes=num_classes, + num_queries=100, + aux_loss=True, + ) + checkpoint = torch.load('./detr-r50-e632da11.pth') + model.load_state_dict(checkpoint['model']) + model.to(device) + model.eval() + + gen_wts(model, "detr") + + # test + # with torch.no_grad(): + # transform = T.Compose([T.Resize(800), T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]) + # im = Image.open('./image/demo.jpg') + # img = transform(im).unsqueeze(0) + + # img = img.to(device) + # res = model(img) + + # logits = res['pred_logits'] + # pred_boxes = res['pred_boxes'] + # out_prob = logits.softmax(-1)[0, :, :-1] + # keep = out_prob.max(-1).values > 0.5 + # label = out_prob[keep].argmax(dim=1) + # out_bbox = pred_boxes[0, keep] + # out_bbox = out_bbox.to(torch.device('cpu')) + # out_bbox = box_cxcywh_to_xyxy(out_bbox) + # out_bbox = out_bbox * torch.tensor([640, 480, 640, 480]) + # image = cv2.imread('./image/demo.jpg') + # for ob in out_bbox: + # x0 = int(ob[0].item()) + # y0 = int(ob[1].item()) + # x1 = int(ob[2].item()) + # y1 = int(ob[3].item()) + # cv2.rectangle(image, (x0, y0), (x1, y1), (0,0,255), 1) + + # cv2.imwrite('res.jpg', image) + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/detr/logging.h b/detr/logging.h new file mode 100644 index 0000000..c9c9d6c --- /dev/null +++ b/detr/logging.h @@ -0,0 +1,447 @@ +/* + * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef TENSORRT_LOGGING_H +#define TENSORRT_LOGGING_H + +#include "NvInferRuntimeCommon.h" +#include +#include +#include +#include +#include +#include +#include + +using Severity = nvinfer1::ILogger::Severity; + +class LogStreamConsumerBuffer : public std::stringbuf { + public: + LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog) + : mOutput(stream) + , mPrefix(prefix) + , mShouldLog(shouldLog) {} + + LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) + : mOutput(other.mOutput) {} + + ~LogStreamConsumerBuffer() { + // std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence + // std::streambuf::pptr() gives a pointer to the current position of the output sequence + // if the pointer to the beginning is not equal to the pointer to the current position, + // call putOutput() to log the output to the stream + if (pbase() != pptr()) { + putOutput(); + } + } + + // synchronizes the stream buffer and returns 0 on success + // synchronizing the stream buffer consists of inserting the buffer contents into the stream, + // resetting the buffer and flushing the stream + virtual int sync() { + putOutput(); + return 0; + } + + void putOutput() { + if (mShouldLog) { + // prepend timestamp + std::time_t timestamp = std::time(nullptr); + tm* tm_local = std::localtime(×tamp); + std::cout << "["; + std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/"; + std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/"; + std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-"; + std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":"; + std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":"; + std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] "; + // std::stringbuf::str() gets the string contents of the buffer + // insert the buffer contents pre-appended by the appropriate prefix into the stream + mOutput << mPrefix << str(); + // set the buffer to empty + str(""); + // flush the stream + mOutput.flush(); + } + } + + void setShouldLog(bool shouldLog) { + mShouldLog = shouldLog; + } + + private: + std::ostream& mOutput; + std::string mPrefix; + bool mShouldLog; +}; + +//! +//! \class LogStreamConsumerBase +//! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer +//! +class LogStreamConsumerBase { + public: + LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog) + : mBuffer(stream, prefix, shouldLog) {} + + protected: + LogStreamConsumerBuffer mBuffer; +}; + +//! +//! \class LogStreamConsumer +//! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages. +//! Order of base classes is LogStreamConsumerBase and then std::ostream. +//! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field +//! in LogStreamConsumer and then the address of the buffer is passed to std::ostream. +//! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream. +//! Please do not change the order of the parent classes. +//! +class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream { + public: + //! \brief Creates a LogStreamConsumer which logs messages with level severity. + //! Reportable severity determines if the messages are severe enough to be logged. + LogStreamConsumer(Severity reportableSeverity, Severity severity) + : LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity) + , std::ostream(&mBuffer) // links the stream buffer with the stream + , mShouldLog(severity <= reportableSeverity) + , mSeverity(severity) {} + + LogStreamConsumer(LogStreamConsumer&& other) + : LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog) + , std::ostream(&mBuffer) // links the stream buffer with the stream + , mShouldLog(other.mShouldLog) + , mSeverity(other.mSeverity) {} + + void setReportableSeverity(Severity reportableSeverity) { + mShouldLog = mSeverity <= reportableSeverity; + mBuffer.setShouldLog(mShouldLog); + } + + private: + static std::ostream& severityOstream(Severity severity) { + return severity >= Severity::kINFO ? std::cout : std::cerr; + } + + static std::string severityPrefix(Severity severity) { + switch (severity) { + case Severity::kINTERNAL_ERROR: return "[F] "; + case Severity::kERROR: return "[E] "; + case Severity::kWARNING: return "[W] "; + case Severity::kINFO: return "[I] "; + case Severity::kVERBOSE: return "[V] "; + default: assert(0); return ""; + } + } + + bool mShouldLog; + Severity mSeverity; +}; + +//! \class Logger +//! +//! \brief Class which manages logging of TensorRT tools and samples +//! +//! \details This class provides a common interface for TensorRT tools and samples to log information to the console, +//! and supports logging two types of messages: +//! +//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal) +//! - Test pass/fail messages +//! +//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is +//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location. +//! +//! In the future, this class could be extended to support dumping test results to a file in some standard format +//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run). +//! +//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger +//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT +//! library and messages coming from the sample. +//! +//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the +//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger +//! object. + +class Logger : public nvinfer1::ILogger { + public: + explicit Logger(Severity severity = Severity::kWARNING) + : mReportableSeverity(severity) {} + + //! + //! \enum TestResult + //! \brief Represents the state of a given test + //! + enum class TestResult { + kRUNNING, //!< The test is running + kPASSED, //!< The test passed + kFAILED, //!< The test failed + kWAIVED //!< The test was waived + }; + + //! + //! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger + //! \return The nvinfer1::ILogger associated with this Logger + //! + //! TODO Once all samples are updated to use this method to register the logger with TensorRT, + //! we can eliminate the inheritance of Logger from ILogger + //! + nvinfer1::ILogger& getTRTLogger() { + return *this; + } + + //! + //! \brief Implementation of the nvinfer1::ILogger::log() virtual method + //! + //! Note samples should not be calling this function directly; it will eventually go away once we eliminate the + //! inheritance from nvinfer1::ILogger + //! + void log(Severity severity, const char* msg) override { + LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl; + } + + //! + //! \brief Method for controlling the verbosity of logging output + //! + //! \param severity The logger will only emit messages that have severity of this level or higher. + //! + void setReportableSeverity(Severity severity) { + mReportableSeverity = severity; + } + + //! + //! \brief Opaque handle that holds logging information for a particular test + //! + //! This object is an opaque handle to information used by the Logger to print test results. + //! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used + //! with Logger::reportTest{Start,End}(). + //! + class TestAtom { + public: + TestAtom(TestAtom&&) = default; + + private: + friend class Logger; + + TestAtom(bool started, const std::string& name, const std::string& cmdline) + : mStarted(started) + , mName(name) + , mCmdline(cmdline) {} + + bool mStarted; + std::string mName; + std::string mCmdline; + }; + + //! + //! \brief Define a test for logging + //! + //! \param[in] name The name of the test. This should be a string starting with + //! "TensorRT" and containing dot-separated strings containing + //! the characters [A-Za-z0-9_]. + //! For example, "TensorRT.sample_googlenet" + //! \param[in] cmdline The command line used to reproduce the test + // + //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). + //! + static TestAtom defineTest(const std::string& name, const std::string& cmdline) { + return TestAtom(false, name, cmdline); + } + + //! + //! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments + //! as input + //! + //! \param[in] name The name of the test + //! \param[in] argc The number of command-line arguments + //! \param[in] argv The array of command-line arguments (given as C strings) + //! + //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). + static TestAtom defineTest(const std::string& name, int argc, char const* const* argv) { + auto cmdline = genCmdlineString(argc, argv); + return defineTest(name, cmdline); + } + + //! + //! \brief Report that a test has started. + //! + //! \pre reportTestStart() has not been called yet for the given testAtom + //! + //! \param[in] testAtom The handle to the test that has started + //! + static void reportTestStart(TestAtom& testAtom) { + reportTestResult(testAtom, TestResult::kRUNNING); + assert(!testAtom.mStarted); + testAtom.mStarted = true; + } + + //! + //! \brief Report that a test has ended. + //! + //! \pre reportTestStart() has been called for the given testAtom + //! + //! \param[in] testAtom The handle to the test that has ended + //! \param[in] result The result of the test. Should be one of TestResult::kPASSED, + //! TestResult::kFAILED, TestResult::kWAIVED + //! + static void reportTestEnd(const TestAtom& testAtom, TestResult result) { + assert(result != TestResult::kRUNNING); + assert(testAtom.mStarted); + reportTestResult(testAtom, result); + } + + static int reportPass(const TestAtom& testAtom) { + reportTestEnd(testAtom, TestResult::kPASSED); + return EXIT_SUCCESS; + } + + static int reportFail(const TestAtom& testAtom) { + reportTestEnd(testAtom, TestResult::kFAILED); + return EXIT_FAILURE; + } + + static int reportWaive(const TestAtom& testAtom) { + reportTestEnd(testAtom, TestResult::kWAIVED); + return EXIT_SUCCESS; + } + + static int reportTest(const TestAtom& testAtom, bool pass) { + return pass ? reportPass(testAtom) : reportFail(testAtom); + } + + Severity getReportableSeverity() const { + return mReportableSeverity; + } + + private: + //! + //! \brief returns an appropriate string for prefixing a log message with the given severity + //! + static const char* severityPrefix(Severity severity) { + switch (severity) { + case Severity::kINTERNAL_ERROR: return "[F] "; + case Severity::kERROR: return "[E] "; + case Severity::kWARNING: return "[W] "; + case Severity::kINFO: return "[I] "; + case Severity::kVERBOSE: return "[V] "; + default: assert(0); return ""; + } + } + + //! + //! \brief returns an appropriate string for prefixing a test result message with the given result + //! + static const char* testResultString(TestResult result) { + switch (result) { + case TestResult::kRUNNING: return "RUNNING"; + case TestResult::kPASSED: return "PASSED"; + case TestResult::kFAILED: return "FAILED"; + case TestResult::kWAIVED: return "WAIVED"; + default: assert(0); return ""; + } + } + + //! + //! \brief returns an appropriate output stream (cout or cerr) to use with the given severity + //! + static std::ostream& severityOstream(Severity severity) { + return severity >= Severity::kINFO ? std::cout : std::cerr; + } + + //! + //! \brief method that implements logging test results + //! + static void reportTestResult(const TestAtom& testAtom, TestResult result) { + severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # " + << testAtom.mCmdline << std::endl; + } + + //! + //! \brief generate a command line string from the given (argc, argv) values + //! + static std::string genCmdlineString(int argc, char const* const* argv) { + std::stringstream ss; + for (int i = 0; i < argc; i++) { + if (i > 0) + ss << " "; + ss << argv[i]; + } + return ss.str(); + } + + Severity mReportableSeverity; +}; + +namespace { + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE +//! +//! Example usage: +//! +//! LOG_VERBOSE(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) { + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO +//! +//! Example usage: +//! +//! LOG_INFO(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_INFO(const Logger& logger) { + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING +//! +//! Example usage: +//! +//! LOG_WARN(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_WARN(const Logger& logger) { + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR +//! +//! Example usage: +//! +//! LOG_ERROR(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_ERROR(const Logger& logger) { + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR +// ("fatal" severity) +//! +//! Example usage: +//! +//! LOG_FATAL(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_FATAL(const Logger& logger) { + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR); +} + +} // anonymous namespace + +#endif // TENSORRT_LOGGING_H