Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Feature/yolov4 p5 #68

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 11 additions & 8 deletions layers/yololayer.cu
Original file line number Diff line number Diff line change
Expand Up @@ -27,14 +27,17 @@ namespace nvinfer1
{
mYoloWidth = yolo_width;
mYoloHeight = yolo_height;
mNumAnchors = num_anchors;
mNumAnchors = num_anchors ;
memcpy(mAnchorsHost, anchors, num_anchors * 2 * sizeof(float));
mNumClasses = num_classes;
mInputWidth = input_width;
mInputHeight = input_height;
mScaleXY = scale_x_y;
mNewCoords = new_coords;
// std::cout << "anchors " << mNumAnchors << " " << MAX_ANCHORS*2<<std::endl;

// for(int i=0; i<num_anchors*2;i++)
// std::cout << anchors[i] << std::endl;
CHECK(cudaMalloc(&mAnchors, MAX_ANCHORS * 2 * sizeof(float)));
CHECK(cudaMemcpy(mAnchors, mAnchorsHost, mNumAnchors * 2 * sizeof(float), cudaMemcpyHostToDevice));
}
Expand All @@ -53,20 +56,19 @@ namespace nvinfer1
read(d, mInputHeight);
read(d, mScaleXY);
read(d, mNewCoords);

CHECK(cudaMalloc(&mAnchors, MAX_ANCHORS * 2 * sizeof(float)));
CHECK(cudaMemcpy(mAnchors, mAnchorsHost, mNumAnchors * 2 * sizeof(float), cudaMemcpyHostToDevice));

assert(d == reinterpret_cast<const char *>(data) + length);
}

void YoloLayerPlugin::serialize(void* buffer) const TRT_NOEXCEPT
{
char* d = static_cast<char*>(buffer);
write(d, mThreadCount);
write(d, mYoloWidth);
write(d, mYoloHeight);
write(d, mNumAnchors);

memcpy(d, mAnchorsHost, MAX_ANCHORS * 2 * sizeof(float));
d += MAX_ANCHORS * 2 * sizeof(float);
write(d, mNumClasses);
Expand Down Expand Up @@ -272,17 +274,18 @@ namespace nvinfer1
class_id = i - 5;
}
}
float box_prob = *(cur_input + 4 * total_grids);
float box_prob = sigmoidGPU(*(cur_input + 4 * total_grids));
//if (max_cls_prob < IGNORE_THRESH || box_prob < IGNORE_THRESH)
// return;

max_cls_prob = sigmoidGPU(max_cls_prob);
int row = (idx % total_grids) / yolo_width;
int col = (idx % total_grids) % yolo_width;

det->bbox[0] = (col + scale(*(cur_input + 0 * total_grids), scale_x_y)) / yolo_width; // [0, 1]
det->bbox[1] = (row + scale(*(cur_input + 1 * total_grids), scale_x_y)) / yolo_height; // [0, 1]
det->bbox[2] = square(*(cur_input + 2 * total_grids)) * 4 * *(anchors + 2 * anchor_idx + 0) / input_w; // [0, 1]
det->bbox[3] = square(*(cur_input + 3 * total_grids)) * 4 * *(anchors + 2 * anchor_idx + 1) / input_h; // [0, 1]
det->bbox[0] = (col + scale_sigmoidGPU(*(cur_input + 0 * total_grids), scale_x_y)) / yolo_width; // [0, 1]
det->bbox[1] = (row + scale_sigmoidGPU(*(cur_input + 1 * total_grids), scale_x_y)) / yolo_height; // [0, 1]
det->bbox[2] = square(sigmoidGPU(*(cur_input + 2 * total_grids))) * 4 * *(anchors + 2 * anchor_idx + 0) / input_w; // [0, 1]
det->bbox[3] = square(sigmoidGPU(*(cur_input + 3 * total_grids))) * 4 * *(anchors + 2 * anchor_idx + 1) / input_h; // [0, 1]

det->bbox[0] -= det->bbox[2] / 2; // shift from center to top-left
det->bbox[1] -= det->bbox[3] / 2;
Expand Down
3 changes: 2 additions & 1 deletion layers/yololayer.h
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,8 @@
#include "math_constants.h"
#include "NvInfer.h"

#define MAX_ANCHORS 6
#define MAX_ANCHORS 8
//TODO make dynamic

#define CHECK(status) \
do { \
Expand Down
258 changes: 161 additions & 97 deletions main.cpp
Original file line number Diff line number Diff line change
@@ -1,114 +1,178 @@
#include "NvInfer.h"
#include "cuda_runtime_api.h"

#include <string>
#include "parser/cxxopts.hpp"


#include "networks/yolov4p5.h"
#include "networks/yolov4.h"
#include "networks/yolov4tiny.h"
#include "networks/yolov4tiny3l.h"

// Don't remove unused include, necessary to correctly load and register tensorrt yolo plugin
#include "layers/yololayer.h"

// Don't remove unused include, necessary to correctly load and register tensorrt yolo plugin
#include "utils/logging.h"


static Logger gLogger;

#include <iostream>

#define DEVICE 0
#define BATCH_SIZE 1

using namespace nvinfer1;

enum NETWORKS {
YOLOV4,
YOLOV4TINY,
YOLOV4TINY3L
};

int main(int argc, char** argv) {
cxxopts::Options options(argv[0], "--- YOLOV4 TRITON TENSORRT ---");

options.add_options()
("n,network", "Network to optimize, either \"yolov4\", \"yolov4tiny\" or \"yolov4tiny3l\"", cxxopts::value<std::string>()->default_value("yolov4"))
("h,help", "Print help screen");

NETWORKS network;

// Parse and check options
try {
auto result = options.parse(argc, argv);

if (result.count("help")) {
std::cout << options.help() << std::endl;
exit(0);
}

auto network_string = result["network"].as<std::string>();
if (network_string.compare("yolov4") == 0) {
network = NETWORKS::YOLOV4;
}
else if (network_string.compare("yolov4tiny") == 0) {
network = NETWORKS::YOLOV4TINY;
}
else if (network_string.compare("yolov4tiny3l") == 0) {
network = NETWORKS::YOLOV4TINY3L;
}
else {
std::cout << "[Error] Network to optimize must be either \"yolov4\" or \"yolov4tiny\"" << std::endl;
std::cout << options.help({""}) << std::endl;
exit(0);
}
// #define BATCH_SIZE 1

using namespace nvinfer1;

enum NETWORKS { YOLOV4, YOLOV4P5, YOLOV4TINY, YOLOV4TINY3L };

int main(int argc, char **argv) {
cxxopts::Options options(argv[0], "--- YOLOV4 TRITON TENSORRT ---");

options.add_options()
("n,network",
"Network to optimize, either \"yolov4\", \"yolov5-p5 \", "
"\"yolov4tiny\" or \"yolov4tiny3l\"",
cxxopts::value<std::string>()->default_value("yolov4"))
("w,weights",
"Path to network weights",
cxxopts::value<std::string>()->default_value("idontexist"))
("b,batch_size",
"Batch size to be used",
cxxopts::value<int>()->default_value("1"))
("input_h",
"Input height",

cxxopts::value<int>()->default_value("608"))
("input_w",
"Input weight",
cxxopts::value<int>()->default_value("608"))
("h,help", "Print help screen");

NETWORKS network;
int BATCH_SIZE = 1;
int INPUT_H = 0;
int INPUT_W = 0;
std::string WEIGHTS_PATH = "idontexist";
// Parse and check options
try {
auto result = options.parse(argc, argv);

if (result.count("help")) {

std::cout << options.help() << std::endl;
exit(0);
}
catch(cxxopts::OptionException exception) {
std::cout << "[Error] " << exception.what() << std::endl;
std::cout << options.help() << std::endl;
exit(0);
WEIGHTS_PATH= result["weights"].as<std::string>();
BATCH_SIZE = result["batch_size"].as<int>();
INPUT_H = result["input_h"].as<int>();
INPUT_W = result["input_w"].as<int>();
auto network_string = result["network"].as<std::string>();
if (network_string.compare("yolov4") == 0) {
network = NETWORKS::YOLOV4;
}
else if (network_string.compare("yolov4p5") == 0) {
network = NETWORKS::YOLOV4P5;
}

cudaSetDevice(DEVICE);

std::cout << "[Info] Creating builder" << std::endl;
// 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;
if(network == NETWORKS::YOLOV4) {
std::cout << "[Info] Creating model yolov4" << std::endl;
engine = yolov4::createEngine(BATCH_SIZE, builder, config, DataType::kFLOAT, "yolov4.wts");
}
else if(network == NETWORKS::YOLOV4TINY) {
std::cout << "[Info] Creating model yolov4tiny" << std::endl;
engine = yolov4tiny::createEngine(BATCH_SIZE, builder, config, DataType::kFLOAT, "yolov4tiny.wts");
else if (network_string.compare("yolov4tiny") == 0) {
network = NETWORKS::YOLOV4TINY;
} else if (network_string.compare("yolov4tiny3l") == 0) {
network = NETWORKS::YOLOV4TINY3L;
} else {
std::cout << "[Error] Network to optimize must be either \"yolov4\" , "
"\"yolov4tiny\" , \"yolov4tiny3l \" or \"yolov4p5\" "
<< std::endl;
std::cout << options.help({""}) << std::endl;
exit(0);
}
else if(network == NETWORKS::YOLOV4TINY3L) {
std::cout << "[Info] Creating model yolov4tiny3l" << std::endl;
engine = yolov4tiny3l::createEngine(BATCH_SIZE, builder, config, DataType::kFLOAT, "yolov4tiny3l.wts");
}
assert(engine != nullptr);

std::cout << "[Info] Serializing model to engine file" << std::endl;
// Serialize the engine
IHostMemory* modelStream{nullptr};
modelStream = engine->serialize();

// Close everything down
engine->destroy();
builder->destroy();

assert(modelStream != nullptr);
std::string engine_name = network == NETWORKS::YOLOV4TINY3L ? "yolov4tiny3l.engine" : ( network == NETWORKS::YOLOV4TINY ? "yolov4tiny.engine" : "yolov4.engine" );
std::ofstream p(engine_name.c_str(), std::ios::binary);
if (!p) {
std::cerr << "[Error] Could not open engine output file " << engine_name << std::endl;
return -1;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
modelStream->destroy();

std::cout << "[Info] Done" << std::endl;

return 0;
}
} catch (cxxopts::OptionException exception) {
std::cout << "[Error] " << exception.what() << std::endl;
std::cout << options.help() << std::endl;
exit(0);
}

cudaSetDevice(DEVICE);

std::cout << "[Info] Creating builder" << std::endl;
// Create builder
IBuilder *builder = createInferBuilder(gLogger);
IBuilderConfig *config = builder->createBuilderConfig();
config->setMaxWorkspaceSize(5*1024*1024);
// Create model to populate the network, then set the outputs and create an
// engine
float gd = 1.0f;
float gw = 1.0f;
ICudaEngine *engine;
if (network == NETWORKS::YOLOV4) {
std::cout << "[Info] Creating model yolov4" << std::endl;
auto params = yolov4::yolov4Parameters();
params.WEIGHTS_PATH= WEIGHTS_PATH;
params.BATCH_SIZE = BATCH_SIZE;
params.INPUT_H = INPUT_H;
params.INPUT_W = INPUT_W;
engine = yolov4::createEngine(params, builder, config, DataType::kFLOAT);
} else if (network == NETWORKS::YOLOV4P5) {
std::cout << "[Info] Creating model yolov4p5" << std::endl;
auto params = yolov4p5::yolov4p5Parameters();
params.WEIGHTS_PATH= WEIGHTS_PATH;
params.BATCH_SIZE = BATCH_SIZE;
params.INPUT_H = INPUT_H;
params.INPUT_W = INPUT_W;
engine = yolov4p5::createEngine(params, builder, config,
DataType::kFLOAT, gd,gw);
} else if (network == NETWORKS::YOLOV4TINY) {
std::cout << "[Info] Creating model yolov4tiny" << std::endl;
// engine = yolov4tiny::createEngine(BATCH_SIZE, builder, config,
// DataType::kFLOAT, "yolov4tiny.wts");
//
auto params = yolov4tiny::yolov4tinyParameters();
params.WEIGHTS_PATH= WEIGHTS_PATH;
params.BATCH_SIZE = BATCH_SIZE;
params.INPUT_H = INPUT_H;
params.INPUT_W = INPUT_W;
engine = yolov4tiny::createEngine(params, builder, config,DataType::kFLOAT);
} else if (network == NETWORKS::YOLOV4TINY3L) {
std::cout << "[Info] Creating model yolov4tiny3l" << std::endl;
// engine = yolov4tiny3l::createEngine(BATCH_SIZE, builder, config,
// DataType::kFLOAT, "yolov4tiny3l.wts");
auto params = yolov4tiny3l::yolov4tiny3lParameters();
params.WEIGHTS_PATH= WEIGHTS_PATH;
params.BATCH_SIZE = BATCH_SIZE;
params.INPUT_H = INPUT_H;
params.INPUT_W = INPUT_W;
engine = yolov4tiny3l::createEngine(params, builder, config,DataType::kFLOAT);
}
assert(engine != nullptr);

std::cout << "[Info] Serializing model to engine file" << std::endl;
// Serialize the engine
IHostMemory *modelStream{nullptr};
modelStream = engine->serialize();

// Close everything down
engine->destroy();
builder->destroy();

assert(modelStream != nullptr);
std::string engine_name = "";
if (network == NETWORKS::YOLOV4TINY3L) {
engine_name = "yolov4tiny3l.engine";
} else if (network == NETWORKS::YOLOV4TINY) {
engine_name = "yolov4tiny.engine";
} else if (network == NETWORKS::YOLOV4) {
engine_name = "yolov4.engine";
} else if (network == NETWORKS::YOLOV4P5) {
engine_name = "yolov4p5.engine";
}
std::ofstream p(engine_name.c_str(), std::ios::binary);
if (!p) {
std::cerr << "[Error] Could not open engine output file " << engine_name
<< std::endl;
return -1;
}
p.write(reinterpret_cast<const char *>(modelStream->data()),
modelStream->size());
modelStream->destroy();

std::cout << "[Info] Done" << std::endl;

return 0;
}
Loading