From 6792b82f4b024389f75fa2283c47d9b15ca9cec2 Mon Sep 17 00:00:00 2001
From: qiuyunzhe <qiuyunzhe94@gmail.com>
Date: Thu, 18 Jun 2020 16:52:08 -0700
Subject: [PATCH] prepare anchor during initialization

---
 yolov4/yololayer.cu | 25 +++++++++++++++++++------
 yolov4/yololayer.h  |  1 +
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/yolov4/yololayer.cu b/yolov4/yololayer.cu
index fd0af16..bbab043 100644
--- a/yolov4/yololayer.cu
+++ b/yolov4/yololayer.cu
@@ -13,6 +13,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()
@@ -32,6 +41,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);
     }
 
@@ -179,9 +197,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);
 
@@ -195,12 +210,10 @@ namespace nvinfer1
             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));
     }
 
 
diff --git a/yolov4/yololayer.h b/yolov4/yololayer.h
index 5074636..af30e9c 100644
--- a/yolov4/yololayer.h
+++ b/yolov4/yololayer.h
@@ -116,6 +116,7 @@ namespace nvinfer1
             int mKernelCount;
             std::vector<Yolo::YoloKernel> mYoloKernel;
             int mThreadCount = 256;
+            void** mAnchor;
             const char* mPluginNamespace;
     };