Fix ONNX Runtime Op memory leak (#2406)

open-mmlab · Sep 6, 2023 · a3c2fae · a3c2fae
1 parent 6fdf459
commit a3c2fae
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/csrc/mmdeploy/backend_ops/onnxruntime/modulated_deform_conv/modulated_deform_conv.cpp b/csrc/mmdeploy/backend_ops/onnxruntime/modulated_deform_conv/modulated_deform_conv.cpp
@@ -190,6 +190,8 @@ void MMCVModulatedDeformConvKernel::Compute(OrtKernelContext *context) {
                              group, deformable_group, channels, num_output, kernel_height,
                              kernel_width, stride_height, stride_width, padding_height,
                              padding_width, dilation_height, dilation_width, columns, out_ptr);
+
+  allocator_.Free(columns);
 }
 REGISTER_ONNXRUNTIME_OPS(mmdeploy, MMCVModulatedDeformConvOp);
 REGISTER_ONNXRUNTIME_OPS(mmcv, MMCVModulatedDeformConvOp);

diff --git a/csrc/mmdeploy/backend_ops/onnxruntime/nms_rotated/nms_rotated.cpp b/csrc/mmdeploy/backend_ops/onnxruntime/nms_rotated/nms_rotated.cpp
@@ -358,6 +358,10 @@ void NMSRotatedKernel::Compute(OrtKernelContext* context) {
   int64_t* res_data = ort_.GetTensorMutableData<int64_t>(res);
 
   memcpy(res_data, res_order.data(), sizeof(int64_t) * res_order.size());
+
+  allocator_.Free(tmp_boxes);
+  allocator_.Free(sc);
+  allocator_.Free(select);
 }
 
 REGISTER_ONNXRUNTIME_OPS(mmdeploy, NMSRotatedOp);