haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add detr (#606)

Browse Source 
* add detr

* Update README.md
This commit is contained in:
JumpPandaer 2021-07-01 13:09:51 +08:00 committed by GitHub
parent 7b85cfd392
commit 161e6b05ee
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 1903 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
detr/CMakeLists.txt Normal file
View File

@ -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)

58
detr/README.md Normal file
View File

@ -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 |

326
detr/backbone.hpp Normal file
View File

@ -0,0 +1,326 @@
#pragma once
#include <map>
#include "common.hpp"
enum RESNETTYPE {
R18 = 0,
R34,
R50,
R101,
R152
};
const std::map<RESNETTYPE, std::vector<int>> 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<std::string, Weights>& 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<float*>(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<float*>(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<float*>(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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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;
}

90
detr/common.hpp Normal file
View File

@ -0,0 +1,90 @@
#pragma once
#include <dirent.h>
#include <cuda_runtime_api.h>
#include <fstream>
#include <sstream>
#include <iostream>
#include <vector>
#include <unordered_map>
#include <algorithm>
#include "./logging.h"
#include <NvInfer.h>
#include <opencv2/opencv.hpp>
static Logger gLogger;
using namespace nvinfer1;
void loadWeights(const std::string file, std::unordered_map<std::string, Weights>& 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<uint32_t*>(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<std::string> &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

826
detr/detr.cpp Normal file
View File

@ -0,0 +1,826 @@
#pragma once
#include <iostream>
#include <unordered_map>
#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<std::string> OUTPUT_NAMES = { "scores", "boxes"};
void preprocessImg(cv::Mat& img) {
// convert to rgb
cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
float ratio = static_cast<float>(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<std::string, Weights>& 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<std::vector<float>> y_embed(h);
for (int i = 0; i < h; i++)
y_embed[i] = std::vector<float>(w, i + 1);
std::vector<float> sub_embed(w, 0);
for (int i = 0; i < w; i++)
sub_embed[i] = i + 1;
std::vector<std::vector<float>> 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<float> 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<float>(num_pos_feats)));
}
// pos_x, pos_y
std::vector<std::vector<std::vector<float>>> pos_x(h,
std::vector<std::vector<float>>(w,
std::vector<float>(num_pos_feats, 0)));
std::vector<std::vector<std::vector<float>>> pos_y(h,
std::vector<std::vector<float>>(w,
std::vector<float>(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<float*>(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<std::string, Weights>& 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<std::string, Weights>& 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<float*>(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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<float*>(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<std::string, Weights>& 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<ITensor*> Predict(
INetworkDefinition *network,
std::unordered_map<std::string, Weights>& 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<ITensor*> 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<std::string, Weights> 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<ITensor*> 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<void*>& buffers,
std::vector<float>& input, std::vector<float*>& 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<const char*>(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<char*>(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<std::string> 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<float> 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<float> scores_h(BATCH_SIZE * NUM_QUERIES * NUM_CLASS);
std::vector<float> boxes_h(BATCH_SIZE * NUM_QUERIES * 4);
std::vector<void*> buffers = { data_d, scores_d, boxes_d };
std::vector<float*> 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<cv::Vec3f>(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<std::chrono::milliseconds>(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;
}

123
detr/gen_wts.py Normal file
View File

@ -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()

447
detr/logging.h Normal file
View File

@ -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 <cassert>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <sstream>
#include <string>
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(&timestamp);
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