diff --git a/yolov4/CMakeLists.txt b/yolov4/CMakeLists.txt
index 36024c1..da27c3a 100644
--- a/yolov4/CMakeLists.txt
+++ b/yolov4/CMakeLists.txt
@@ -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")
diff --git a/yolov4/README.md b/yolov4/README.md
index 946f114..552233b 100644
--- a/yolov4/README.md
+++ b/yolov4/README.md
@@ -2,48 +2,6 @@
The Pytorch implementation is from [ultralytics/yolov3](https://github.com/ultralytics/yolov3). It can load yolov4.cfg and yolov4.weights(from AlexeyAB/darknet).
-Following tricks are used in this yolov4:
-
-- Three yololayer are implemented in one plugin to improve speed, codes derived from [lewes6369/TensorRT-Yolov3](https://github.com/lewes6369/TensorRT-Yolov3)
-- Mish activation, implemented in a plugin.
-- Batchnorm layer, implemented by scale layer.
-
-## Excute:
-
-```
-1. generate yolov4.wts from pytorch implementation with yolov4.cfg and yolov4.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 yolov4.weights from https://github.com/AlexeyAB/darknet#pre-trained-models
-cd yolov3
-cp ../tensorrtx/yolov4/gen_wts.py .
-python gen_wts.py yolov4.weights
-// a file 'yolov4.wts' will be generated.
-// the master branch of yolov3 should work, if not, you can checkout be87b41aa2fe59be8e62f4b488052b24ad0bd450
-
-2. put yolov4.wts into ./yolov4, build and run
-
-mv yolov4.wts ../tensorrtx/yolov4/
-cd ../tensorrtx/yolov4
-mkdir build
-cd build
-cmake ..
-make
-sudo ./yolov4 -s // serialize model to plan file i.e. 'yolov4.engine'
-sudo ./yolov4 -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
-```
-
-
-
-
-
-
-
-
-
## Config
- Input shape `INPUT_H`, `INPUT_W` defined in yololayer.h
@@ -54,6 +12,44 @@ sudo ./yolov4 -d ../../yolov3-spp/samples // deserialize plan file and run infe
- bbox confidence threshold `BBOX_CONF_THRESH` in yolov4.cpp
- `BATCH_SIZE` in yolov4.cpp
+## How to run
+
+1. generate yolov4.wts from pytorch implementation with yolov4.cfg and yolov4.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 yolov4.weights from https://github.com/AlexeyAB/darknet#pre-trained-models
+cp {tensorrtx}/yolov4/gen_wts.py {ultralytics/yolov3/}
+cd {ultralytics/yolov3/}
+python gen_wts.py yolov4.weights
+// a file 'yolov4.wts' will be generated.
+// the master branch of yolov3 should work, if not, you can checkout be87b41aa2fe59be8e62f4b488052b24ad0bd450
+```
+
+2. put yolov4.wts into {tensorrtx}/yolov4, build and run
+
+```
+mv yolov4.wts {tensorrtx}/yolov4/
+cd {tensorrtx}/yolov4
+mkdir build
+cd build
+cmake ..
+make
+sudo ./yolov4 -s // serialize model to plan file i.e. 'yolov4.engine'
+sudo ./yolov4 -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
+
+
+
+
+
+
+
+
+
## More Information
See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx)
diff --git a/yolov4/Utils.h b/yolov4/Utils.h
deleted file mode 100644
index 0de663c..0000000
--- a/yolov4/Utils.h
+++ /dev/null
@@ -1,94 +0,0 @@
-#ifndef __TRT_UTILS_H_
-#define __TRT_UTILS_H_
-
-#include
-#include
-#include
-#include
-
-#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 Record;
- std::vector 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
- void write(char*& buffer, const T& val)
- {
- *reinterpret_cast(buffer) = val;
- buffer += sizeof(T);
- }
-
- template
- void read(const char*& buffer, T& val)
- {
- val = *reinterpret_cast(buffer);
- buffer += sizeof(T);
- }
-}
-
-#endif
\ No newline at end of file
diff --git a/yolov4/utils.h b/yolov4/utils.h
new file mode 100644
index 0000000..09981ff
--- /dev/null
+++ b/yolov4/utils.h
@@ -0,0 +1,39 @@
+#ifndef __TRT_UTILS_H_
+#define __TRT_UTILS_H_
+
+#include
+#include
+#include
+#include
+
+#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
+ void write(char*& buffer, const T& val)
+ {
+ *reinterpret_cast(buffer) = val;
+ buffer += sizeof(T);
+ }
+
+ template
+ void read(const char*& buffer, T& val)
+ {
+ val = *reinterpret_cast(buffer);
+ buffer += sizeof(T);
+ }
+}
+
+#endif
diff --git a/yolov4/yololayer.cu b/yolov4/yololayer.cu
index bbab043..3344483 100644
--- a/yolov4/yololayer.cu
+++ b/yolov4/yololayer.cu
@@ -1,4 +1,6 @@
+#include
#include "yololayer.h"
+#include "utils.h"
using namespace Yolo;
diff --git a/yolov4/yololayer.h b/yolov4/yololayer.h
index 47d73ee..c7bea1c 100644
--- a/yolov4/yololayer.h
+++ b/yolov4/yololayer.h
@@ -1,13 +1,9 @@
#ifndef _YOLO_LAYER_H
#define _YOLO_LAYER_H
-#include
-#include
-#include
-#include
-#include "NvInfer.h"
-#include "Utils.h"
#include
+#include
+#include "NvInfer.h"
namespace Yolo
{
diff --git a/yolov4/yolov4.cpp b/yolov4/yolov4.cpp
index 0e7e566..9e23baf 100644
--- a/yolov4/yolov4.cpp
+++ b/yolov4/yolov4.cpp
@@ -7,23 +7,12 @@
#include
#include
#include "NvInfer.h"
-#include "NvInferPlugin.h"
+#include "utils.h"
#include "cuda_runtime_api.h"
#include "logging.h"
#include "yololayer.h"
#include "mish.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 DEVICE 0 // GPU id
#define NMS_THRESH 0.4
@@ -57,7 +46,7 @@ cv::Mat preprocess_img(cv::Mat& img) {
y = 0;
}
cv::Mat re(h, w, CV_8UC3);
- cv::resize(img, re, re.size(), 0, 0, cv::INTER_CUBIC);
+ cv::resize(img, re, re.size());
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;
@@ -180,7 +169,6 @@ IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
@@ -209,7 +197,6 @@ IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int p, int linx) {
- std::cout << linx << std::endl;
Weights emptywts{DataType::kFLOAT, nullptr, 0};
IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap["module_list." + std::to_string(linx) + ".Conv2d.weight"], emptywts);
assert(conv1);
@@ -227,7 +214,6 @@ ILayer* convBnMish(INetworkDefinition *network, std::map&
}
ILayer* convBnLeaky(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int p, int linx) {
- std::cout << linx << std::endl;
Weights emptywts{DataType::kFLOAT, nullptr, 0};
IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap["module_list." + std::to_string(linx) + ".Conv2d.weight"], emptywts);
assert(conv1);
@@ -489,7 +475,6 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilder
auto yolo = network->addPluginV2(inputTensors_yolo, 3, *pluginObj);
yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME);
- std::cout << "set name out" << std::endl;
network->markOutput(*yolo->getOutput(0));
// Build engine
@@ -498,8 +483,9 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilder
#ifdef USE_FP16
config->setFlag(BuilderFlag::kFP16);
#endif
+ std::cout << "Building tensorrt engine, please wait for a while..." << std::endl;
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
- std::cout << "build out" << std::endl;
+ std::cout << "Build engine successfully!" << std::endl;
// Don't need the network any more
network->destroy();
@@ -528,6 +514,7 @@ void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) {
// Close everything down
engine->destroy();
builder->destroy();
+ config->destroy();
}
void doInference(IExecutionContext& context, float* input, float* output, int batchSize) {
@@ -544,23 +531,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]));
+ CUDA_CHECK(cudaFree(buffers[inputIndex]));
+ CUDA_CHECK(cudaFree(buffers[outputIndex]));
}
int read_files_in_dir(const char *p_dir_name, std::vector &file_names) {