thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update yoloPlugins.cpp

Browse Source
This commit is contained in:
Marcos Luciano
2022-02-18 03:51:08 -03:00
committed by GitHub
parent 329daeb150
commit 0c815aef80
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
nvdsinfer_custom_impl_Yolo/yoloPlugins.cpp
View File

@@ -212,8 +212,8 @@ int YoloLayer::enqueue(
} }
else { else {
void* softmax; void* softmax;
cudaMallocHost(&softmax, sizeof(outputs[0])); CHECK(cudaMallocHost(&softmax, sizeof(outputs[0])));
cudaMemcpy(softmax, outputs[0], sizeof(outputs[0]), cudaMemcpyHostToDevice); CHECK(cudaMemcpy(softmax, outputs[0], sizeof(outputs[0]), cudaMemcpyHostToDevice));
CHECK(cudaYoloLayer_v2( CHECK(cudaYoloLayer_v2(
inputs[0], outputs[0], softmax, batchSize, m_GridSizeX, m_GridSizeY, m_NumClasses, m_NumBoxes, inputs[0], outputs[0], softmax, batchSize, m_GridSizeX, m_GridSizeY, m_NumClasses, m_NumBoxes,
@@ -276,4 +276,4 @@ nvinfer1::IPluginV2* YoloLayer::clone() const noexcept
return new YoloLayer (m_NumBoxes, m_NumClasses, m_GridSizeX, m_GridSizeY, m_Type, m_NewCoords, m_ScaleXY, m_BetaNMS, m_Anchors, m_Mask); return new YoloLayer (m_NumBoxes, m_NumClasses, m_GridSizeX, m_GridSizeY, m_Type, m_NewCoords, m_ScaleXY, m_BetaNMS, m_Anchors, m_Mask);
} }
REGISTER_TENSORRT_PLUGIN(YoloLayerPluginCreator); REGISTER_TENSORRT_PLUGIN(YoloLayerPluginCreator);