yolov5s batchsize
This commit is contained in:
parent
6ee4c62ff0
commit
271b152c30
@ -49,6 +49,7 @@ Following models are implemented.
|
||||
|[yolov3](./yolov3)| darknet-53, weights and pytorch implementation from [ultralytics/yolov3](https://github.com/ultralytics/yolov3) |
|
||||
|[yolov3-spp](./yolov3-spp)| darknet-53, weights and pytorch implementation from [ultralytics/yolov3](https://github.com/ultralytics/yolov3) |
|
||||
|[yolov4](./yolov4)| CSPDarknet53, weights from [AlexeyAB/darknet](https://github.com/AlexeyAB/darknet#pre-trained-models), pytorch implementation from [ultralytics/yolov3](https://github.com/ultralytics/yolov3) |
|
||||
|[yolov5](./yolov5)| yolov5-s, pytorch implementation from [ultralytics/yolov5](https://github.com/ultralytics/yolov5) |
|
||||
|[retinaface](./retinaface)| resnet-50, weights from [biubug6/Pytorch_Retinaface](https://github.com/biubug6/Pytorch_Retinaface) |
|
||||
|[arcface](./arcface)| LResNet50E-IR, weights from [deepinsight/insightface](https://github.com/deepinsight/insightface) |
|
||||
|[retinafaceAntiCov](./retinafaceAntiCov)| mobilenet0.25, weights from [deepinsight/insightface](https://github.com/deepinsight/insightface), retinaface anti-COVID-19, detect face and mask attribute |
|
||||
@ -85,6 +86,9 @@ Some tricky operations encountered in these models, already solved, but might ha
|
||||
| YOLOv4(CSPDarknet53) | Xeon E5-2620/GTX1080 | 1 | FP16 | 608x608 | 35.7 |
|
||||
| YOLOv4(CSPDarknet53) | Xeon E5-2620/GTX1080 | 4 | FP16 | 608x608 | 40.9 |
|
||||
| YOLOv4(CSPDarknet53) | Xeon E5-2620/GTX1080 | 8 | FP16 | 608x608 | 41.3 |
|
||||
| YOLOv5-s | Xeon E5-2620/GTX1080 | 1 | FP16 | 608x608 | 167 |
|
||||
| YOLOv5-s | Xeon E5-2620/GTX1080 | 4 | FP16 | 608x608 | 182 |
|
||||
| YOLOv5-s | Xeon E5-2620/GTX1080 | 8 | FP16 | 608x608 | 186 |
|
||||
| RetinaFace(resnet50) | TX2 | 1 | FP16 | 384x640 | 15 |
|
||||
| RetinaFace(resnet50) | Xeon E5-2620/GTX1080 | 1 | FP32 | 928x1600 | 15 |
|
||||
| ArcFace(LResNet50E-IR) | Xeon E5-2620/GTX1080 | 1 | FP32 | 112x112 | 333 |
|
||||
|
||||
@ -45,6 +45,7 @@ sudo ./yolov5s -d ../samples // deserialize plan file and run inference, the im
|
||||
- GPU id can be selected by the macro in yolov5s.cpp
|
||||
- NMS thresh in yolov5s.cpp
|
||||
- BBox confidence thresh in yolov5s.cpp
|
||||
- Batch size in yolov5s.cpp
|
||||
|
||||
## More Information
|
||||
|
||||
|
||||
@ -15,6 +15,15 @@ namespace nvinfer1
|
||||
mYoloKernel.push_back(yolo3);
|
||||
|
||||
mKernelCount = mYoloKernel.size();
|
||||
|
||||
CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*)));
|
||||
size_t AnchorLen = sizeof(float)* CHECK_COUNT*2;
|
||||
for(int ii = 0; ii < mKernelCount; ii ++)
|
||||
{
|
||||
CUDA_CHECK(cudaMalloc(&mAnchor[ii],AnchorLen));
|
||||
const auto& yolo = mYoloKernel[ii];
|
||||
CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice));
|
||||
}
|
||||
}
|
||||
|
||||
YoloLayerPlugin::~YoloLayerPlugin()
|
||||
@ -34,6 +43,15 @@ namespace nvinfer1
|
||||
memcpy(mYoloKernel.data(),d,kernelSize);
|
||||
d += kernelSize;
|
||||
|
||||
CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*)));
|
||||
size_t AnchorLen = sizeof(float)* CHECK_COUNT*2;
|
||||
for(int ii = 0; ii < mKernelCount; ii ++)
|
||||
{
|
||||
CUDA_CHECK(cudaMalloc(&mAnchor[ii],AnchorLen));
|
||||
const auto& yolo = mYoloKernel[ii];
|
||||
CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice));
|
||||
}
|
||||
|
||||
assert(d == a + length);
|
||||
}
|
||||
|
||||
@ -162,7 +180,7 @@ namespace nvinfer1
|
||||
float *res_count = output + bnIdx*outputElem;
|
||||
int count = (int)atomicAdd(res_count, 1);
|
||||
if (count >= MAX_OUTPUT_BBOX_COUNT) return;
|
||||
char* data = (char * )res_count + sizeof(float) + count*sizeof(Detection);
|
||||
char* data = (char *)res_count + sizeof(float) + count * sizeof(Detection);
|
||||
Detection* det = (Detection*)(data);
|
||||
|
||||
int row = idx / yoloWidth;
|
||||
@ -181,9 +199,6 @@ namespace nvinfer1
|
||||
}
|
||||
|
||||
void YoloLayerPlugin::forwardGpu(const float *const * inputs, float* output, cudaStream_t stream, int batchSize) {
|
||||
void* devAnchor;
|
||||
size_t AnchorLen = sizeof(float)* CHECK_COUNT*2;
|
||||
CUDA_CHECK(cudaMalloc(&devAnchor,AnchorLen));
|
||||
|
||||
int outputElem = 1 + MAX_OUTPUT_BBOX_COUNT * sizeof(Detection) / sizeof(float);
|
||||
|
||||
@ -191,28 +206,22 @@ namespace nvinfer1
|
||||
CUDA_CHECK(cudaMemset(output + idx*outputElem, 0, sizeof(float)));
|
||||
}
|
||||
int numElem = 0;
|
||||
for (unsigned int i = 0;i< mYoloKernel.size();++i)
|
||||
for (unsigned int i = 0; i < mYoloKernel.size(); ++i)
|
||||
{
|
||||
const auto& yolo = mYoloKernel[i];
|
||||
numElem = yolo.width*yolo.height*batchSize;
|
||||
if (numElem < mThreadCount)
|
||||
mThreadCount = numElem;
|
||||
CUDA_CHECK(cudaMemcpy(devAnchor, yolo.anchors, AnchorLen, cudaMemcpyHostToDevice));
|
||||
CalDetection<<< (yolo.width*yolo.height*batchSize + mThreadCount - 1) / mThreadCount, mThreadCount>>>
|
||||
(inputs[i],output, numElem, yolo.width, yolo.height, (float *)devAnchor, mClassCount ,outputElem);
|
||||
(inputs[i], output, numElem, yolo.width, yolo.height, (float *)mAnchor[i], mClassCount, outputElem);
|
||||
}
|
||||
|
||||
CUDA_CHECK(cudaFree(devAnchor));
|
||||
}
|
||||
|
||||
|
||||
int YoloLayerPlugin::enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream)
|
||||
{
|
||||
//assert(batchSize == 1);
|
||||
//GPU
|
||||
//CUDA_CHECK(cudaStreamSynchronize(stream));
|
||||
forwardGpu((const float *const *)inputs, (float*)outputs[0], stream, batchSize);
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
||||
@ -110,6 +110,7 @@ namespace nvinfer1
|
||||
int mKernelCount;
|
||||
std::vector<Yolo::YoloKernel> mYoloKernel;
|
||||
int mThreadCount = 256;
|
||||
void** mAnchor;
|
||||
const char* mPluginNamespace;
|
||||
};
|
||||
|
||||
|
||||
@ -8,6 +8,7 @@
|
||||
#define DEVICE 0 // GPU id
|
||||
#define NMS_THRESH 0.5
|
||||
#define CONF_THRESH 0.4
|
||||
#define BATCH_SIZE 1
|
||||
|
||||
// stuff we know about the network and the input/output blobs
|
||||
static const int INPUT_H = Yolo::INPUT_H;
|
||||
@ -159,7 +160,7 @@ int main(int argc, char** argv) {
|
||||
|
||||
if (argc == 2 && std::string(argv[1]) == "-s") {
|
||||
IHostMemory* modelStream{nullptr};
|
||||
APIToModel(1, &modelStream);
|
||||
APIToModel(BATCH_SIZE, &modelStream);
|
||||
assert(modelStream != nullptr);
|
||||
std::ofstream p("yolov5s.engine", std::ios::binary);
|
||||
if (!p) {
|
||||
@ -194,10 +195,10 @@ int main(int argc, char** argv) {
|
||||
}
|
||||
|
||||
// prepare input data ---------------------------
|
||||
float data[3 * INPUT_H * INPUT_W];
|
||||
static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W];
|
||||
//for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++)
|
||||
// data[i] = 1.0;
|
||||
static float prob[OUTPUT_SIZE];
|
||||
static float prob[BATCH_SIZE * OUTPUT_SIZE];
|
||||
IRuntime* runtime = createInferRuntime(gLogger);
|
||||
assert(runtime != nullptr);
|
||||
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
|
||||
@ -207,40 +208,42 @@ int main(int argc, char** argv) {
|
||||
delete[] trtModelStream;
|
||||
|
||||
int fcount = 0;
|
||||
for (auto f: file_names) {
|
||||
for (int f = 0; f < (int)file_names.size(); f++) {
|
||||
fcount++;
|
||||
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);
|
||||
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;
|
||||
data[i + 2 * INPUT_H * INPUT_W] = pr_img.at<cv::Vec3b>(i)[0] / 255.0;
|
||||
if (fcount < BATCH_SIZE && f + 1 != (int)file_names.size()) continue;
|
||||
for (int b = 0; b < fcount; b++) {
|
||||
cv::Mat img = cv::imread(std::string(argv[2]) + "/" + file_names[f - fcount + 1 + b]);
|
||||
if (img.empty()) continue;
|
||||
cv::Mat pr_img = preprocess_img(img);
|
||||
for (int i = 0; i < INPUT_H * INPUT_W; i++) {
|
||||
data[b * 3 * INPUT_H * INPUT_W + i] = pr_img.at<cv::Vec3b>(i)[2] / 255.0;
|
||||
data[b * 3 * INPUT_H * INPUT_W + i + INPUT_H * INPUT_W] = pr_img.at<cv::Vec3b>(i)[1] / 255.0;
|
||||
data[b * 3 * INPUT_H * INPUT_W + i + 2 * INPUT_H * INPUT_W] = pr_img.at<cv::Vec3b>(i)[0] / 255.0;
|
||||
}
|
||||
}
|
||||
|
||||
// Run inference
|
||||
auto start = std::chrono::system_clock::now();
|
||||
doInference(*context, data, prob, 1);
|
||||
doInference(*context, data, prob, BATCH_SIZE);
|
||||
auto end = std::chrono::system_clock::now();
|
||||
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
|
||||
std::vector<Yolo::Detection> res;
|
||||
nms(res, prob, CONF_THRESH, NMS_THRESH);
|
||||
for (int i=0; i<20; i++) {
|
||||
std::cout << prob[i] << ",";
|
||||
std::vector<std::vector<Yolo::Detection>> batch_res(fcount);
|
||||
for (int b = 0; b < fcount; b++) {
|
||||
auto& res = batch_res[b];
|
||||
nms(res, &prob[b * OUTPUT_SIZE], CONF_THRESH, NMS_THRESH);
|
||||
}
|
||||
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] << ", ";
|
||||
for (int b = 0; b < fcount; b++) {
|
||||
auto& res = batch_res[b];
|
||||
//std::cout << res.size() << std::endl;
|
||||
cv::Mat img = cv::imread(std::string(argv[2]) + "/" + file_names[f - fcount + 1 + b]);
|
||||
for (size_t j = 0; j < res.size(); j++) {
|
||||
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);
|
||||
}
|
||||
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);
|
||||
cv::imwrite("_" + file_names[f - fcount + 1 + b], img);
|
||||
}
|
||||
cv::imwrite("_" + f, img);
|
||||
fcount = 0;
|
||||
}
|
||||
|
||||
// Destroy the engine
|
||||
|
||||
Loading…
Reference in New Issue
Block a user