yolov3 support int8
This commit is contained in:
parent
5b94c930f8
commit
d16bb0c624
@ -113,6 +113,7 @@ Some tricky operations encountered in these models, already solved, but might ha
|
||||
|-|-|:-:|:-:|:-:|:-:|
|
||||
| YOLOv3-tiny | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 333 |
|
||||
| YOLOv3(darknet53) | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 39.2 |
|
||||
| YOLOv3(darknet53) | Xeon E5-2620/GTX1080 | 1 | INT8 | 608x608 | 71.4 |
|
||||
| YOLOv3-spp(darknet53) | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 38.5 |
|
||||
| YOLOv4(CSPDarknet53) | Xeon E5-2620/GTX1080 | 1 | FP32 | 608x608 | 35.7 |
|
||||
| YOLOv4(CSPDarknet53) | Xeon E5-2620/GTX1080 | 4 | FP32 | 608x608 | 40.9 |
|
||||
|
||||
@ -13,16 +13,13 @@ find_package(CUDA REQUIRED)
|
||||
set(CUDA_NVCC_PLAGS ${CUDA_NVCC_PLAGS};-std=c++11;-g;-G;-gencode;arch=compute_30;code=sm_30)
|
||||
|
||||
include_directories(${PROJECT_SOURCE_DIR}/include)
|
||||
if (CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64")
|
||||
message("embed_platform on")
|
||||
include_directories(/usr/local/cuda/targets/aarch64-linux/include)
|
||||
link_directories(/usr/local/cuda/targets/aarch64-linux/lib)
|
||||
else()
|
||||
message("embed_platform off")
|
||||
include_directories(/usr/local/cuda/include)
|
||||
link_directories(/usr/local/cuda/lib64)
|
||||
endif()
|
||||
|
||||
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
|
||||
# cuda
|
||||
include_directories(/usr/local/cuda/include)
|
||||
link_directories(/usr/local/cuda/lib64)
|
||||
# tensorrt
|
||||
include_directories(/usr/include/x86_64-linux-gnu/)
|
||||
link_directories(/usr/lib/x86_64-linux-gnu/)
|
||||
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
|
||||
|
||||
@ -33,7 +30,7 @@ target_link_libraries(yololayer nvinfer cudart)
|
||||
find_package(OpenCV)
|
||||
include_directories(OpenCV_INCLUDE_DIRS)
|
||||
|
||||
add_executable(yolov3 ${PROJECT_SOURCE_DIR}/yolov3.cpp)
|
||||
add_executable(yolov3 ${PROJECT_SOURCE_DIR}/calibrator.cpp ${PROJECT_SOURCE_DIR}/yolov3.cpp)
|
||||
target_link_libraries(yolov3 nvinfer)
|
||||
target_link_libraries(yolov3 cudart)
|
||||
target_link_libraries(yolov3 yololayer)
|
||||
|
||||
@ -4,33 +4,54 @@ The Pytorch implementation is [ultralytics/yolov3](https://github.com/ultralytic
|
||||
|
||||
This branch is using tensorrt7 API, there is also a yolov3 implementation using tensorrt4 API, go to [branch trt4/yolov3](https://github.com/wang-xinyu/tensorrtx/tree/trt4/yolov3), which is using [ayooshkathuria/pytorch-yolo-v3](https://github.com/ayooshkathuria/pytorch-yolo-v3).
|
||||
|
||||
## Excute:
|
||||
## Config
|
||||
|
||||
- Input shape defined in yololayer.h
|
||||
- Number of classes defined in yololayer.h
|
||||
- INT8/FP16/FP32 can be selected by the macro in yolov3.cpp
|
||||
- GPU id can be selected by the macro in yolov3.cpp
|
||||
- NMS thresh in yolov3.cpp
|
||||
- BBox confidence thresh in yolov3.cpp
|
||||
|
||||
## How to run
|
||||
|
||||
```
|
||||
1. generate yolov3.wts from pytorch implementation with yolov3.cfg and yolov3.weights, or download .wts from model zoo
|
||||
|
||||
```
|
||||
git clone https://github.com/wang-xinyu/tensorrtx.git
|
||||
git clone https://github.com/ultralytics/yolov3.git
|
||||
// download its weights 'yolov3.pt' or 'yolov3.weights'
|
||||
cd yolov3
|
||||
cp ../tensorrtx/yolov3/gen_wts.py .
|
||||
cp {tensorrtx}/yolov3/gen_wts.py {ultralytics/yolov3/}
|
||||
cd {ultralytics/yolov3/}
|
||||
python gen_wts.py yolov3.weights
|
||||
// a file 'yolov3.wts' will be generated.
|
||||
// the master branch of yolov3 should work, if not, you can checkout cf7a4d31d37788023a9186a1a143a2dab0275ead
|
||||
```
|
||||
|
||||
2. put yolov3.wts into tensorrtx/yolov3, build and run
|
||||
|
||||
mv yolov3.wts ../tensorrtx/yolov3/
|
||||
cd ../tensorrtx/yolov3
|
||||
```
|
||||
mv yolov3.wts {tensorrtx}/yolov3/
|
||||
cd {tensorrtx}/yolov3
|
||||
mkdir build
|
||||
cd build
|
||||
cmake ..
|
||||
make
|
||||
sudo ./yolov3 -s // serialize model to plan file i.e. 'yolov3.engine'
|
||||
sudo ./yolov3 -d ../../yolov3-spp/samples // deserialize plan file and run inference, the images in samples will be processed.
|
||||
sudo ./yolov3 -s // serialize model to plan file i.e. 'yolov3.engine'
|
||||
sudo ./yolov3 -d ../../yolov3-spp/samples // deserialize plan file and run inference, the images in samples will be processed.
|
||||
```
|
||||
|
||||
3. check the images generated, as follows. _zidane.jpg and _bus.jpg
|
||||
```
|
||||
|
||||
# INT8 Quantization
|
||||
|
||||
1. Prepare calibration images, you can randomly select 1000s images from your train set. For coco, you can also download my calibration images `coco_calib` from [BaiduPan](https://pan.baidu.com/s/1GOm_-JobpyLMAqZWCDUhKg) pwd: a9wh
|
||||
|
||||
2. unzip it in yolov3/build
|
||||
|
||||
3. set the macro `USE_INT8` in yolov3.cpp and make
|
||||
|
||||
4. serialize the model and test
|
||||
|
||||
<p align="center">
|
||||
<img src="https://user-images.githubusercontent.com/15235574/78247927-4d9fac00-751e-11ea-8b1b-704a0aeb3fcf.jpg">
|
||||
@ -40,15 +61,6 @@ sudo ./yolov3 -d ../../yolov3-spp/samples // deserialize plan file and run infe
|
||||
<img src="https://user-images.githubusercontent.com/15235574/78247970-60b27c00-751e-11ea-88df-41473fed4823.jpg">
|
||||
</p>
|
||||
|
||||
## Config
|
||||
|
||||
- Input shape defined in yololayer.h
|
||||
- Number of classes defined in yololayer.h
|
||||
- FP16/FP32 can be selected by the macro in yolov3.cpp
|
||||
- GPU id can be selected by the macro in yolov3.cpp
|
||||
- NMS thresh in yolov3.cpp
|
||||
- BBox confidence thresh in yolov3.cpp
|
||||
|
||||
## More Information
|
||||
|
||||
See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx)
|
||||
|
||||
@ -1,94 +0,0 @@
|
||||
#ifndef __TRT_UTILS_H_
|
||||
#define __TRT_UTILS_H_
|
||||
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <algorithm>
|
||||
#include <cudnn.h>
|
||||
|
||||
#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
|
||||
|
||||
namespace Tn
|
||||
{
|
||||
class Profiler : public nvinfer1::IProfiler
|
||||
{
|
||||
public:
|
||||
void printLayerTimes(int itrationsTimes)
|
||||
{
|
||||
float totalTime = 0;
|
||||
for (size_t i = 0; i < mProfile.size(); i++)
|
||||
{
|
||||
printf("%-40.40s %4.3fms\n", mProfile[i].first.c_str(), mProfile[i].second / itrationsTimes);
|
||||
totalTime += mProfile[i].second;
|
||||
}
|
||||
printf("Time over all layers: %4.3f\n", totalTime / itrationsTimes);
|
||||
}
|
||||
private:
|
||||
typedef std::pair<std::string, float> Record;
|
||||
std::vector<Record> mProfile;
|
||||
|
||||
virtual void reportLayerTime(const char* layerName, float ms)
|
||||
{
|
||||
auto record = std::find_if(mProfile.begin(), mProfile.end(), [&](const Record& r){ return r.first == layerName; });
|
||||
if (record == mProfile.end())
|
||||
mProfile.push_back(std::make_pair(layerName, ms));
|
||||
else
|
||||
record->second += ms;
|
||||
}
|
||||
};
|
||||
|
||||
//Logger for TensorRT info/warning/errors
|
||||
class Logger : public nvinfer1::ILogger
|
||||
{
|
||||
public:
|
||||
|
||||
Logger(): Logger(Severity::kWARNING) {}
|
||||
|
||||
Logger(Severity severity): reportableSeverity(severity) {}
|
||||
|
||||
void log(Severity severity, const char* msg) override
|
||||
{
|
||||
// suppress messages with severity enum value greater than the reportable
|
||||
if (severity > reportableSeverity) return;
|
||||
|
||||
switch (severity)
|
||||
{
|
||||
case Severity::kINTERNAL_ERROR: std::cerr << "INTERNAL_ERROR: "; break;
|
||||
case Severity::kERROR: std::cerr << "ERROR: "; break;
|
||||
case Severity::kWARNING: std::cerr << "WARNING: "; break;
|
||||
case Severity::kINFO: std::cerr << "INFO: "; break;
|
||||
default: std::cerr << "UNKNOWN: "; break;
|
||||
}
|
||||
std::cerr << msg << std::endl;
|
||||
}
|
||||
|
||||
Severity reportableSeverity{Severity::kWARNING};
|
||||
};
|
||||
|
||||
template<typename T>
|
||||
void write(char*& buffer, const T& val)
|
||||
{
|
||||
*reinterpret_cast<T*>(buffer) = val;
|
||||
buffer += sizeof(T);
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
void read(const char*& buffer, T& val)
|
||||
{
|
||||
val = *reinterpret_cast<const T*>(buffer);
|
||||
buffer += sizeof(T);
|
||||
}
|
||||
}
|
||||
|
||||
#endif
|
||||
80
yolov3/calibrator.cpp
Normal file
80
yolov3/calibrator.cpp
Normal file
@ -0,0 +1,80 @@
|
||||
#include <iostream>
|
||||
#include <iterator>
|
||||
#include <fstream>
|
||||
#include <opencv2/dnn/dnn.hpp>
|
||||
#include "calibrator.h"
|
||||
#include "cuda_runtime_api.h"
|
||||
#include "utils.h"
|
||||
|
||||
Int8EntropyCalibrator2::Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache)
|
||||
: batchsize_(batchsize)
|
||||
, input_w_(input_w)
|
||||
, input_h_(input_h)
|
||||
, img_idx_(0)
|
||||
, img_dir_(img_dir)
|
||||
, calib_table_name_(calib_table_name)
|
||||
, input_blob_name_(input_blob_name)
|
||||
, read_cache_(read_cache)
|
||||
{
|
||||
input_count_ = 3 * input_w * input_h * batchsize;
|
||||
CUDA_CHECK(cudaMalloc(&device_input_, input_count_ * sizeof(float)));
|
||||
read_files_in_dir(img_dir, img_files_);
|
||||
}
|
||||
|
||||
Int8EntropyCalibrator2::~Int8EntropyCalibrator2()
|
||||
{
|
||||
CUDA_CHECK(cudaFree(device_input_));
|
||||
}
|
||||
|
||||
int Int8EntropyCalibrator2::getBatchSize() const
|
||||
{
|
||||
return batchsize_;
|
||||
}
|
||||
|
||||
bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int nbBindings)
|
||||
{
|
||||
if (img_idx_ + batchsize_ > (int)img_files_.size()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
std::vector<cv::Mat> input_imgs_;
|
||||
for (int i = img_idx_; i < img_idx_ + batchsize_; i++) {
|
||||
std::cout << img_files_[i] << " " << i << std::endl;
|
||||
cv::Mat temp = cv::imread(img_dir_ + img_files_[i]);
|
||||
if (temp.empty()){
|
||||
std::cerr << "Fatal error: image cannot open!" << std::endl;
|
||||
return false;
|
||||
}
|
||||
cv::Mat pr_img = preprocess_img(temp, input_w_, input_h_);
|
||||
input_imgs_.push_back(pr_img);
|
||||
}
|
||||
img_idx_ += batchsize_;
|
||||
cv::Mat blob = cv::dnn::blobFromImages(input_imgs_, 1.0 / 255.0, cv::Size(input_w_, input_h_), cv::Scalar(0, 0, 0), true, false);
|
||||
|
||||
CUDA_CHECK(cudaMemcpy(device_input_, blob.ptr<float>(0), input_count_ * sizeof(float), cudaMemcpyHostToDevice));
|
||||
assert(!strcmp(names[0], input_blob_name_));
|
||||
bindings[0] = device_input_;
|
||||
return true;
|
||||
}
|
||||
|
||||
const void* Int8EntropyCalibrator2::readCalibrationCache(size_t& length)
|
||||
{
|
||||
std::cout << "reading calib cache: " << calib_table_name_ << std::endl;
|
||||
calib_cache_.clear();
|
||||
std::ifstream input(calib_table_name_, std::ios::binary);
|
||||
input >> std::noskipws;
|
||||
if (read_cache_ && input.good())
|
||||
{
|
||||
std::copy(std::istream_iterator<char>(input), std::istream_iterator<char>(), std::back_inserter(calib_cache_));
|
||||
}
|
||||
length = calib_cache_.size();
|
||||
return length ? calib_cache_.data() : nullptr;
|
||||
}
|
||||
|
||||
void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, size_t length)
|
||||
{
|
||||
std::cout << "writing calib cache: " << calib_table_name_ << " size: " << length << std::endl;
|
||||
std::ofstream output(calib_table_name_, std::ios::binary);
|
||||
output.write(reinterpret_cast<const char*>(cache), length);
|
||||
}
|
||||
|
||||
39
yolov3/calibrator.h
Normal file
39
yolov3/calibrator.h
Normal file
@ -0,0 +1,39 @@
|
||||
#ifndef ENTROPY_CALIBRATOR_H
|
||||
#define ENTROPY_CALIBRATOR_H
|
||||
|
||||
#include "NvInfer.h"
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
//! \class Int8EntropyCalibrator2
|
||||
//!
|
||||
//! \brief Implements Entropy calibrator 2.
|
||||
//! CalibrationAlgoType is kENTROPY_CALIBRATION_2.
|
||||
//!
|
||||
class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2
|
||||
{
|
||||
public:
|
||||
Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache = true);
|
||||
|
||||
virtual ~Int8EntropyCalibrator2();
|
||||
int getBatchSize() const override;
|
||||
bool getBatch(void* bindings[], const char* names[], int nbBindings) override;
|
||||
const void* readCalibrationCache(size_t& length) override;
|
||||
void writeCalibrationCache(const void* cache, size_t length) override;
|
||||
|
||||
private:
|
||||
int batchsize_;
|
||||
int input_w_;
|
||||
int input_h_;
|
||||
int img_idx_;
|
||||
std::string img_dir_;
|
||||
std::vector<std::string> img_files_;
|
||||
size_t input_count_;
|
||||
std::string calib_table_name_;
|
||||
const char* input_blob_name_;
|
||||
bool read_cache_;
|
||||
void* device_input_;
|
||||
std::vector<char> calib_cache_;
|
||||
};
|
||||
|
||||
#endif // ENTROPY_CALIBRATOR_H
|
||||
85
yolov3/utils.h
Normal file
85
yolov3/utils.h
Normal file
@ -0,0 +1,85 @@
|
||||
#ifndef __TRT_UTILS_H_
|
||||
#define __TRT_UTILS_H_
|
||||
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <algorithm>
|
||||
#include <cudnn.h>
|
||||
#include <dirent.h>
|
||||
#include <opencv2/opencv.hpp>
|
||||
|
||||
#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
|
||||
|
||||
namespace Tn
|
||||
{
|
||||
template<typename T>
|
||||
void write(char*& buffer, const T& val)
|
||||
{
|
||||
*reinterpret_cast<T*>(buffer) = val;
|
||||
buffer += sizeof(T);
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
void read(const char*& buffer, T& val)
|
||||
{
|
||||
val = *reinterpret_cast<const T*>(buffer);
|
||||
buffer += sizeof(T);
|
||||
}
|
||||
}
|
||||
|
||||
static inline cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h) {
|
||||
int w, h, x, y;
|
||||
float r_w = input_w / (img.cols*1.0);
|
||||
float r_h = input_h / (img.rows*1.0);
|
||||
if (r_h > r_w) {
|
||||
w = input_w;
|
||||
h = r_w * img.rows;
|
||||
x = 0;
|
||||
y = (input_h - h) / 2;
|
||||
} else {
|
||||
w = r_h * img.cols;
|
||||
h = input_h;
|
||||
x = (input_w - w) / 2;
|
||||
y = 0;
|
||||
}
|
||||
cv::Mat re(h, w, CV_8UC3);
|
||||
cv::resize(img, re, re.size(), 0, 0, cv::INTER_LINEAR);
|
||||
cv::Mat out(input_h, input_w, CV_8UC3, cv::Scalar(128, 128, 128));
|
||||
re.copyTo(out(cv::Rect(x, y, re.cols, re.rows)));
|
||||
return out;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
#endif
|
||||
@ -1,4 +1,6 @@
|
||||
#include "yololayer.h"
|
||||
#include "utils.h"
|
||||
#include <assert.h>
|
||||
|
||||
using namespace Yolo;
|
||||
|
||||
|
||||
@ -1,13 +1,9 @@
|
||||
#ifndef _YOLO_LAYER_H
|
||||
#define _YOLO_LAYER_H
|
||||
|
||||
#include <assert.h>
|
||||
#include <cmath>
|
||||
#include <string.h>
|
||||
#include <cublas_v2.h>
|
||||
#include "NvInfer.h"
|
||||
#include "Utils.h"
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include "NvInfer.h"
|
||||
|
||||
namespace Yolo
|
||||
{
|
||||
@ -51,7 +47,6 @@ namespace Yolo
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
class YoloLayerPlugin: public IPluginV2IOExt
|
||||
|
||||
@ -4,26 +4,14 @@
|
||||
#include <sstream>
|
||||
#include <vector>
|
||||
#include <chrono>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <dirent.h>
|
||||
#include "NvInfer.h"
|
||||
#include "cuda_runtime_api.h"
|
||||
#include "utils.h"
|
||||
#include "logging.h"
|
||||
#include "yololayer.h"
|
||||
#include "calibrator.h"
|
||||
|
||||
#define CHECK(status) \
|
||||
do\
|
||||
{\
|
||||
auto ret = (status);\
|
||||
if (ret != 0)\
|
||||
{\
|
||||
std::cerr << "Cuda failure: " << ret << std::endl;\
|
||||
abort();\
|
||||
}\
|
||||
} while (0)
|
||||
|
||||
|
||||
#define USE_FP16 // comment out this if want to use FP32
|
||||
#define USE_FP16 // set USE_INT8 or USE_FP16 or USE_FP32
|
||||
#define DEVICE 0 // GPU id
|
||||
#define NMS_THRESH 0.4
|
||||
#define BBOX_CONF_THRESH 0.5
|
||||
@ -33,33 +21,12 @@ using namespace nvinfer1;
|
||||
// stuff we know about the network and the input/output blobs
|
||||
static const int INPUT_H = Yolo::INPUT_H;
|
||||
static const int INPUT_W = Yolo::INPUT_W;
|
||||
static const int OUTPUT_SIZE = 1000 * 7 + 1; // we assume the yololayer outputs no more than 1000 boxes that conf >= 0.1
|
||||
static const int DETECTION_SIZE = sizeof(Yolo::Detection) / sizeof(float);
|
||||
static const int OUTPUT_SIZE = Yolo::MAX_OUTPUT_BBOX_COUNT * DETECTION_SIZE + 1; // we assume the yololayer outputs no more than MAX_OUTPUT_BBOX_COUNT boxes that conf >= 0.1
|
||||
const char* INPUT_BLOB_NAME = "data";
|
||||
const char* OUTPUT_BLOB_NAME = "prob";
|
||||
static Logger gLogger;
|
||||
|
||||
cv::Mat preprocess_img(cv::Mat& img) {
|
||||
int w, h, x, y;
|
||||
float r_w = INPUT_W / (img.cols*1.0);
|
||||
float r_h = INPUT_H / (img.rows*1.0);
|
||||
if (r_h > r_w) {
|
||||
w = INPUT_W;
|
||||
h = r_w * img.rows;
|
||||
x = 0;
|
||||
y = (INPUT_H - h) / 2;
|
||||
} else {
|
||||
w = r_h* img.cols;
|
||||
h = INPUT_H;
|
||||
x = (INPUT_W - w) / 2;
|
||||
y = 0;
|
||||
}
|
||||
cv::Mat re(h, w, CV_8UC3);
|
||||
cv::resize(img, re, re.size(), 0, 0, cv::INTER_CUBIC);
|
||||
cv::Mat out(INPUT_H, INPUT_W, CV_8UC3, cv::Scalar(128, 128, 128));
|
||||
re.copyTo(out(cv::Rect(x, y, re.cols, re.rows)));
|
||||
return out;
|
||||
}
|
||||
|
||||
cv::Rect get_rect(cv::Mat& img, float bbox[4]) {
|
||||
int l, r, t, b;
|
||||
float r_w = INPUT_W / (img.cols * 1.0);
|
||||
@ -370,9 +337,16 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilder
|
||||
// Build engine
|
||||
builder->setMaxBatchSize(maxBatchSize);
|
||||
config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB
|
||||
#ifdef USE_FP16
|
||||
#if defined(USE_FP16)
|
||||
config->setFlag(BuilderFlag::kFP16);
|
||||
#elif defined(USE_INT8)
|
||||
std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
|
||||
assert(builder->platformHasFastInt8());
|
||||
config->setFlag(BuilderFlag::kINT8);
|
||||
Int8EntropyCalibrator2 *calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME);
|
||||
config->setInt8Calibrator(calibrator);
|
||||
#endif
|
||||
|
||||
std::cout << "Building engine, please wait for a while..." << std::endl;
|
||||
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
|
||||
std::cout << "Build engine successfully!" << std::endl;
|
||||
@ -420,45 +394,23 @@ void doInference(IExecutionContext& context, float* input, float* output, int ba
|
||||
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
|
||||
|
||||
// Create GPU buffers on device
|
||||
CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
|
||||
CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
|
||||
CUDA_CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
|
||||
CUDA_CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
|
||||
|
||||
// Create stream
|
||||
cudaStream_t stream;
|
||||
CHECK(cudaStreamCreate(&stream));
|
||||
CUDA_CHECK(cudaStreamCreate(&stream));
|
||||
|
||||
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
|
||||
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
|
||||
CUDA_CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
|
||||
context.enqueue(batchSize, buffers, stream, nullptr);
|
||||
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
|
||||
CUDA_CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
|
||||
cudaStreamSynchronize(stream);
|
||||
|
||||
// Release stream and buffers
|
||||
cudaStreamDestroy(stream);
|
||||
CHECK(cudaFree(buffers[inputIndex]));
|
||||
CHECK(cudaFree(buffers[outputIndex]));
|
||||
}
|
||||
|
||||
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;
|
||||
CUDA_CHECK(cudaFree(buffers[inputIndex]));
|
||||
CUDA_CHECK(cudaFree(buffers[outputIndex]));
|
||||
}
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
@ -522,7 +474,7 @@ int main(int argc, char** argv) {
|
||||
std::cout << fcount << " " << f << std::endl;
|
||||
cv::Mat img = cv::imread(std::string(argv[2]) + "/" + f);
|
||||
if (img.empty()) continue;
|
||||
cv::Mat pr_img = preprocess_img(img);
|
||||
cv::Mat pr_img = preprocess_img(img, INPUT_W, INPUT_H);
|
||||
for (int i = 0; i < INPUT_H * INPUT_W; i++) {
|
||||
data[i] = pr_img.at<cv::Vec3b>(i)[2] / 255.0;
|
||||
data[i + INPUT_H * INPUT_W] = pr_img.at<cv::Vec3b>(i)[1] / 255.0;
|
||||
@ -536,16 +488,7 @@ int main(int argc, char** argv) {
|
||||
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
|
||||
std::vector<Yolo::Detection> res;
|
||||
nms(res, prob);
|
||||
for (int i=0; i<20; i++) {
|
||||
std::cout << prob[i] << ",";
|
||||
}
|
||||
std::cout << res.size() << std::endl;
|
||||
for (size_t j = 0; j < res.size(); j++) {
|
||||
float *p = (float*)&res[j];
|
||||
for (size_t k = 0; k < 7; k++) {
|
||||
std::cout << p[k] << ", ";
|
||||
}
|
||||
std::cout << std::endl;
|
||||
cv::Rect r = get_rect(img, res[j].bbox);
|
||||
cv::rectangle(img, r, cv::Scalar(0x27, 0xC1, 0x36), 2);
|
||||
cv::putText(img, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
|
||||
|
||||
Loading…
Reference in New Issue
Block a user