Keep double-buffer reader for static mode (PaddlePaddle#49068)

lxsbupt · Dec 14, 2022 · 2fc38e3 · 2fc38e3
1 parent 22ab027
commit 2fc38e3
Showing 1 changed file with 14 additions and 5 deletions.
diff --git a/python/paddle/fluid/reader.py b/python/paddle/fluid/reader.py
@@ -51,6 +51,7 @@
 from .layers.io import (
     monkey_patch_reader_methods,
     _copy_reader_var_,
+    __create_unshared_decorated_reader__,
 )
 from .unique_name import UniqueNameGenerator
 from .framework import _get_paddle_place, _get_paddle_place_list
@@ -1351,11 +1352,6 @@ def __init__(
         self._use_double_buffer = use_double_buffer
         self._capacity = capacity
         if not self._iterable:
-            # Because layers.io.double_buffer is not supported anymore and that iterable is False and use_double_buffer
-            # is True is not spported, here if itrable is False, use_double_buffer will be
-            # forcely set False to avoid unexpected error.
-            # TODO: keep use_double_buffer
-            self._use_double_buffer = False
             self._init_non_iterable()
 
     def _wait_thread_ends(self):
@@ -1410,6 +1406,7 @@ def _init_non_iterable(self):
             'lod_tensor_blocking_queue'
         )
         reader_name = data_loader_unique_name_generator('create_py_reader')
+        double_buffer_name = data_loader_unique_name_generator('double_buffer')
 
         var = global_scope().var(queue_name)
         self._queue = core.init_lod_tensor_blocking_queue(
@@ -1455,6 +1452,18 @@ def _init_non_iterable(self):
 
             reader = monkey_patch_reader_methods(main_prog_var)
 
+        if self._use_double_buffer:
+            double_buffer_reader = __create_unshared_decorated_reader__(
+                'create_double_buffer_reader',
+                reader,
+                {},
+                name=double_buffer_name,
+            )
+            # we return a double buffer reader. However, the reset method comes from
+            # py_reader.
+            double_buffer_reader.reset = reader.reset
+            reader = double_buffer_reader
+
         self._reader = reader
 
         default_main_program().current_block().append_op(