This is an automated email from the git hooks/post-receive script.
Git pushed a commit to branch master
in repository ffmpeg.
The following commit(s) were added to refs/heads/master by this push:
new c60ac8f256 avfilter/dnn: implement batching for Torch backend
c60ac8f256 is described below
commit c60ac8f256cd76341df4fb5abdfd38d9c47c0a6e
Author: Raja-89 <[email protected]>
AuthorDate: Sat Jul 11 22:30:07 2026 +0530
Commit: guoyejun <[email protected]>
CommitDate: Sat Jul 18 02:56:18 2026 +0000
avfilter/dnn: implement batching for Torch backend
Add batch processing support to the LibTorch DNN backend, following
the same pattern used by the OpenVINO backend.
Key changes:
- Add batch_size AVOption (range 1-1000, default 1) to DnnContext
- Accumulate lltasks in the queue, trigger inference when batch_size
lltasks are ready (matching the OpenVINO batch execution pattern)
- Pre-allocate a single contiguous memory buffer for the entire batch
in fill_model_input_th() to avoid unnecessary tensor copies
- Split batched output in infer_completion_callback() and dispatch
each slice to its corresponding task
Tested with:
ffmpeg -f lavfi -i testsrc=duration=5:size=640x480:rate=25 -vf
format=rgb24,dnn_processing=dnn_backend=torch:model=dummy_model.pt:batch_size=4
-f null -
Signed-off-by: Raja Rathour <[email protected]>
---
libavfilter/dnn/dnn_backend_torch.cpp | 197 +++++++++++++++++++++-------------
libavfilter/dnn/dnn_interface.c | 2 +
libavfilter/dnn_interface.h | 1 +
3 files changed, 125 insertions(+), 75 deletions(-)
diff --git a/libavfilter/dnn/dnn_backend_torch.cpp
b/libavfilter/dnn/dnn_backend_torch.cpp
index 2beffd6e2c..9ba6d61377 100644
--- a/libavfilter/dnn/dnn_backend_torch.cpp
+++ b/libavfilter/dnn/dnn_backend_torch.cpp
@@ -52,7 +52,8 @@ typedef struct THInferRequest {
typedef struct THRequestItem {
THInferRequest *infer_request;
- LastLevelTaskItem *lltask;
+ LastLevelTaskItem **lltasks;
+ uint32_t lltask_count;
DNNAsyncExecModule exec_module;
} THRequestItem;
@@ -108,7 +109,7 @@ static inline void destroy_request_item(THRequestItem **arg)
item = *arg;
th_free_request(item->infer_request);
av_freep(&item->infer_request);
- av_freep(&item->lltask);
+ av_freep(&item->lltasks);
ff_dnn_async_module_cleanup(&item->exec_module);
av_freep(arg);
}
@@ -166,53 +167,75 @@ static int fill_model_input_th(THModel *th_model,
THRequestItem *request)
DNNData input = { 0 };
DnnContext *ctx = th_model->ctx;
int ret, width_idx, height_idx, channel_idx;
+ int batch_size = ctx->batch_size;
+ float *batch_data = NULL;
+ int frame_size = 0;
- lltask = (LastLevelTaskItem *)ff_queue_pop_front(th_model->lltask_queue);
- if (!lltask) {
- ret = AVERROR(EINVAL);
- goto err;
- }
- request->lltask = lltask;
- task = lltask->task;
infer_request = request->infer_request;
ret = get_input_th(&th_model->model, &input, NULL);
- if ( ret != 0) {
+ if (ret != 0) {
goto err;
}
width_idx = dnn_get_width_idx_by_layout(input.layout);
height_idx = dnn_get_height_idx_by_layout(input.layout);
channel_idx = dnn_get_channel_idx_by_layout(input.layout);
+
+ lltask = (LastLevelTaskItem *)ff_queue_peek_front(th_model->lltask_queue);
+ if (!lltask) {
+ ret = AVERROR(EINVAL);
+ goto err;
+ }
+ task = lltask->task;
input.dims[height_idx] = task->in_frame->height;
input.dims[width_idx] = task->in_frame->width;
- input.data = av_malloc(input.dims[height_idx] * input.dims[width_idx] *
- input.dims[channel_idx] * sizeof(float));
- if (!input.data)
- return AVERROR(ENOMEM);
- infer_request->input_tensor = new torch::Tensor();
- infer_request->output = new torch::Tensor();
- switch (th_model->model.func_type) {
- case DFT_PROCESS_FRAME:
- input.scale = 255;
- if (task->do_ioproc) {
- if (th_model->model.frame_pre_proc != NULL) {
- th_model->model.frame_pre_proc(task->in_frame, &input,
th_model->model.filter_ctx);
- } else {
- ff_proc_from_frame_to_dnn(task->in_frame, &input, ctx);
+ frame_size = input.dims[height_idx] * input.dims[width_idx] *
input.dims[channel_idx];
+ batch_data = (float *)av_malloc(batch_size * frame_size * sizeof(float));
+ if (!batch_data) {
+ ret = AVERROR(ENOMEM);
+ goto err;
+ }
+
+ for (int i = 0; i < batch_size; i++) {
+ lltask = (LastLevelTaskItem
*)ff_queue_pop_front(th_model->lltask_queue);
+ if (!lltask)
+ break;
+
+ request->lltasks[i] = lltask;
+ request->lltask_count = i + 1;
+ task = lltask->task;
+
+ input.data = batch_data + i * frame_size;
+
+ switch (th_model->model.func_type) {
+ case DFT_PROCESS_FRAME:
+ input.scale = 255;
+ if (task->do_ioproc) {
+ if (th_model->model.frame_pre_proc != NULL) {
+ th_model->model.frame_pre_proc(task->in_frame, &input,
th_model->model.filter_ctx);
+ } else {
+ ff_proc_from_frame_to_dnn(task->in_frame, &input, ctx);
+ }
}
+ break;
+ default:
+ avpriv_report_missing_feature(NULL, "model function type %d",
th_model->model.func_type);
+ break;
}
- break;
- default:
- avpriv_report_missing_feature(NULL, "model function type %d",
th_model->model.func_type);
- break;
}
- *infer_request->input_tensor = torch::from_blob(input.data,
- {1, input.dims[channel_idx], input.dims[height_idx],
input.dims[width_idx]},
+
+ infer_request->input_tensor = new torch::Tensor();
+ infer_request->output = new torch::Tensor();
+ *infer_request->input_tensor = torch::from_blob(batch_data,
+ {request->lltask_count, input.dims[channel_idx],
input.dims[height_idx], input.dims[width_idx]},
deleter, torch::kFloat32);
+
return 0;
err:
+ if (batch_data)
+ av_freep(&batch_data);
th_free_request(infer_request);
return ret;
}
@@ -233,7 +256,7 @@ static int th_start_inference(void *args)
return AVERROR(EINVAL);
}
infer_request = request->infer_request;
- lltask = request->lltask;
+ lltask = request->lltasks[0];
task = lltask->task;
th_model = (THModel *)task->model;
ctx = th_model->ctx;
@@ -260,54 +283,66 @@ static int th_start_inference(void *args)
static void infer_completion_callback(void *args) {
THRequestItem *request = (THRequestItem*)args;
- LastLevelTaskItem *lltask = request->lltask;
- TaskItem *task = lltask->task;
- DNNData outputs = { 0 };
THInferRequest *infer_request = request->infer_request;
- THModel *th_model = (THModel *)task->model;
+ LastLevelTaskItem *lltask = request->lltasks[0];
+ THModel *th_model = (THModel *)lltask->task->model;
torch::Tensor *output = infer_request->output;
+ DNNData outputs = { 0 };
- c10::IntArrayRef sizes = output->sizes();
- outputs.order = DCO_RGB;
- outputs.layout = DL_NCHW;
- outputs.dt = DNN_FLOAT;
- if (sizes.size() == 4) {
- // 4 dimensions: [batch_size, channel, height, width]
- // this format of data is normally used for video frame SR
- outputs.dims[0] = sizes.at(0); // N
- outputs.dims[1] = sizes.at(1); // C
- outputs.dims[2] = sizes.at(2); // H
- outputs.dims[3] = sizes.at(3); // W
- } else {
- avpriv_report_missing_feature(th_model->ctx, "Support of this kind of
model");
- goto err;
- }
+ auto slices = torch::split(*output, /*split_size=*/1, /*dim=*/0);
+ for (uint32_t i = 0; i < request->lltask_count; i++) {
+ lltask = request->lltasks[i];
+ TaskItem *task = lltask->task;
+ torch::Tensor out_slice = slices[i];
+ c10::IntArrayRef sizes = out_slice.sizes();
+
+ outputs.order = DCO_RGB;
+ outputs.layout = DL_NCHW;
+ outputs.dt = DNN_FLOAT;
+
+ if (sizes.size() == 4) {
+ // 4 dimensions: [batch_size, channel, height, width]
+ // this format of data is normally used for video frame SR
+ outputs.dims[0] = sizes.at(0); // N
+ outputs.dims[1] = sizes.at(1); // C
+ outputs.dims[2] = sizes.at(2); // H
+ outputs.dims[3] = sizes.at(3); // W
+ } else {
+ avpriv_report_missing_feature(th_model->ctx, "Support of this kind
of model");
+ goto err;
+ }
- switch (th_model->model.func_type) {
- case DFT_PROCESS_FRAME:
- if (task->do_ioproc) {
- // Post process can only deal with CPU memory.
- if (output->device() != torch::kCPU)
- *output = output->to(torch::kCPU);
- outputs.scale = 255;
- outputs.data = output->data_ptr();
- if (th_model->model.frame_post_proc != NULL) {
- th_model->model.frame_post_proc(task->out_frame, &outputs,
th_model->model.filter_ctx);
+ switch (th_model->model.func_type) {
+ case DFT_PROCESS_FRAME:
+ if (task->do_ioproc) {
+ // Post process can only deal with CPU memory.
+ if (out_slice.device() != torch::kCPU)
+ out_slice = out_slice.to(torch::kCPU);
+ outputs.scale = 255;
+ outputs.data = out_slice.data_ptr();
+ if (th_model->model.frame_post_proc != NULL) {
+ th_model->model.frame_post_proc(task->out_frame, &outputs,
th_model->model.filter_ctx);
+ } else {
+ ff_proc_from_dnn_to_frame(task->out_frame, &outputs,
th_model->ctx);
+ }
} else {
- ff_proc_from_dnn_to_frame(task->out_frame, &outputs,
th_model->ctx);
+ task->out_frame->width =
outputs.dims[dnn_get_width_idx_by_layout(outputs.layout)];
+ task->out_frame->height =
outputs.dims[dnn_get_height_idx_by_layout(outputs.layout)];
}
- } else {
- task->out_frame->width =
outputs.dims[dnn_get_width_idx_by_layout(outputs.layout)];
- task->out_frame->height =
outputs.dims[dnn_get_height_idx_by_layout(outputs.layout)];
+ break;
+ default:
+ avpriv_report_missing_feature(th_model->ctx, "model function type
%d", th_model->model.func_type);
+ goto err;
}
- break;
- default:
- avpriv_report_missing_feature(th_model->ctx, "model function type %d",
th_model->model.func_type);
- goto err;
+ task->inference_done++;
}
- task->inference_done++;
- av_freep(&request->lltask);
+
err:
+ for (uint32_t i = 0; i < request->lltask_count; i++) {
+ av_freep(&request->lltasks[i]);
+ }
+ request->lltask_count = 0;
+
th_free_request(infer_request);
if (ff_safe_queue_push_back(th_model->request_queue, request) < 0) {
@@ -483,6 +518,11 @@ static DNNModel *dnn_load_model_th(DnnContext *ctx,
DNNFunctionType func_type, A
if (!item->infer_request) {
goto fail;
}
+ item->lltasks = (LastLevelTaskItem **)av_malloc_array(ctx->batch_size,
sizeof(*item->lltasks));
+ if (!item->lltasks) {
+ goto fail;
+ }
+ item->lltask_count = 0;
item->exec_module.start_inference = &th_start_inference;
item->exec_module.callback = &infer_completion_callback;
@@ -551,13 +591,20 @@ static int dnn_execute_model_th(const DNNModel *model,
DNNExecBaseParams *exec_p
return ret;
}
- request = (THRequestItem
*)ff_safe_queue_pop_front(th_model->request_queue);
- if (!request) {
- av_log(ctx, AV_LOG_ERROR, "unable to get infer request.\n");
- return AVERROR(EINVAL);
+ while (ff_queue_size(th_model->lltask_queue) >= ctx->batch_size) {
+ request = (THRequestItem
*)ff_safe_queue_pop_front(th_model->request_queue);
+ if (!request) {
+ av_log(ctx, AV_LOG_ERROR, "unable to get infer request.\n");
+ return AVERROR(EINVAL);
+ }
+
+ ret = execute_model_th(request, th_model->lltask_queue);
+ if (ret != 0) {
+ return ret;
+ }
}
- return execute_model_th(request, th_model->lltask_queue);
+ return 0;
}
static DNNAsyncStatusType dnn_get_result_th(const DNNModel *model, AVFrame
**in, AVFrame **out)
diff --git a/libavfilter/dnn/dnn_interface.c b/libavfilter/dnn/dnn_interface.c
index 010677dd81..2cf81793de 100644
--- a/libavfilter/dnn/dnn_interface.c
+++ b/libavfilter/dnn/dnn_interface.c
@@ -52,6 +52,8 @@ static const AVOption dnn_base_options[] = {
OFFSET(backend_options), AV_OPT_TYPE_STRING, {.str = NULL}, 0,
0, FLAGS | AV_OPT_FLAG_DEPRECATED},
{"nireq", "number of request",
OFFSET(nireq), AV_OPT_TYPE_INT, {.i64 = 0}, 0, INT_MAX, FLAGS},
+ {"batch_size", "batch size per request",
+ OFFSET(batch_size), AV_OPT_TYPE_INT, {.i64 = 1}, 1, 1000,
FLAGS},
{"async", "use DNN async inference",
OFFSET(async), AV_OPT_TYPE_BOOL, {.i64 = 1}, 0, 1, FLAGS},
{"device", "device to run model",
diff --git a/libavfilter/dnn_interface.h b/libavfilter/dnn_interface.h
index 69a8b0a669..207a9c93d7 100644
--- a/libavfilter/dnn_interface.h
+++ b/libavfilter/dnn_interface.h
@@ -165,6 +165,7 @@ typedef struct DnnContext {
const DNNModule *dnn_module;
int nireq;
+ int batch_size;
char *device;
int device_id;
_______________________________________________
ffmpeg-cvslog mailing list -- [email protected]
To unsubscribe send an email to [email protected]
