[FFmpeg-devel,5/8] libavfilter/dnn: avoid memcpy for tensorflow dnn output

Submitted by Guo, Yejun on April 2, 2019, 2:29 p.m.

Details

Message ID 1554215349-13962-1-git-send-email-yejun.guo@intel.com
State New
Headers show

Commit Message

Guo, Yejun April 2, 2019, 2:29 p.m.
use TF_Tensor's cpu address to avoid extra memcpy.

Signed-off-by: Guo, Yejun <yejun.guo@intel.com>
---
 libavfilter/dnn_backend_tf.c | 37 ++++++++++++-------------------------
 libavfilter/vf_sr.c          |  3 ---
 2 files changed, 12 insertions(+), 28 deletions(-)

Patch hide | download patch | download mbox

diff --git a/libavfilter/dnn_backend_tf.c b/libavfilter/dnn_backend_tf.c
index 7966688..be8401e 100644
--- a/libavfilter/dnn_backend_tf.c
+++ b/libavfilter/dnn_backend_tf.c
@@ -35,6 +35,7 @@  typedef struct TFModel{
     TF_Status *status;
     TF_Output input, output;
     TF_Tensor *input_tensor;
+    TF_Tensor *output_tensor;
 } TFModel;
 
 static void free_buffer(void *data, size_t length)
@@ -460,14 +461,11 @@  DNNModel *ff_dnn_load_model_tf(const char *model_filename)
         return NULL;
     }
 
-    tf_model = av_malloc(sizeof(TFModel));
+    tf_model = av_mallocz(sizeof(TFModel));
     if (!tf_model){
         av_freep(&model);
         return NULL;
     }
-    tf_model->session = NULL;
-    tf_model->input_tensor = NULL;
-    tf_model->output_data = NULL;
 
     if (load_tf_model(tf_model, model_filename) != DNN_SUCCESS){
         if (load_native_model(tf_model, model_filename) != DNN_SUCCESS){
@@ -489,36 +487,22 @@  DNNModel *ff_dnn_load_model_tf(const char *model_filename)
 DNNReturnType ff_dnn_execute_model_tf(const DNNModel *model, DNNData *output)
 {
     TFModel *tf_model = (TFModel *)model->model;
-    TF_Tensor *output_tensor;
-    uint32_t count;
-    uint32_t old_count = output->height * output->width * output->channels * sizeof(float);
+    if (tf_model->output_tensor)
+        TF_DeleteTensor(tf_model->output_tensor);
 
     TF_SessionRun(tf_model->session, NULL,
                   &tf_model->input, &tf_model->input_tensor, 1,
-                  &tf_model->output, &output_tensor, 1,
+                  &tf_model->output, &tf_model->output_tensor, 1,
                   NULL, 0, NULL, tf_model->status);
 
     if (TF_GetCode(tf_model->status) != TF_OK){
         return DNN_ERROR;
     }
 
-    output->height = TF_Dim(output_tensor, 1);
-    output->width = TF_Dim(output_tensor, 2);
-    output->channels = TF_Dim(output_tensor, 3);
-    count = output->height * output->width * output->channels * sizeof(float);
-    if (output->data) {
-        if (count > old_count) {
-            av_freep(&output->data);
-        }
-    }
-    if (!output->data) {
-        output->data = av_malloc(count);
-        if (!output->data){
-            return DNN_ERROR;
-        }
-    }
-    memcpy(output->data, TF_TensorData(output_tensor), count);
-    TF_DeleteTensor(output_tensor);
+    output->height = TF_Dim(tf_model->output_tensor, 1);
+    output->width = TF_Dim(tf_model->output_tensor, 2);
+    output->channels = TF_Dim(tf_model->output_tensor, 3);
+    output->data = TF_TensorData(tf_model->output_tensor);
 
     return DNN_SUCCESS;
 }
@@ -542,6 +526,9 @@  void ff_dnn_free_model_tf(DNNModel **model)
         if (tf_model->input_tensor){
             TF_DeleteTensor(tf_model->input_tensor);
         }
+        if (tf_model->output_tensor){
+            TF_DeleteTensor(tf_model->output_tensor);
+        }
         av_freep(&tf_model);
         av_freep(model);
     }
diff --git a/libavfilter/vf_sr.c b/libavfilter/vf_sr.c
index 7c92730..53bd8ea 100644
--- a/libavfilter/vf_sr.c
+++ b/libavfilter/vf_sr.c
@@ -277,9 +277,6 @@  static av_cold void uninit(AVFilterContext *context)
     int i;
     SRContext *sr_context = context->priv;
 
-    if (sr_context->backend_type == DNN_TF)
-        av_freep(&sr_context->output.data);
-
     if (sr_context->dnn_module){
         (sr_context->dnn_module->free_model)(&sr_context->model);
         av_freep(&sr_context->dnn_module);