Skip to content

Commit b048ddf

Browse files
committed
Merge error
1 parent b8975d6 commit b048ddf

File tree

1 file changed

+12
-4
lines changed

1 file changed

+12
-4
lines changed

paddle/fluid/operators/reader/open_files_op.cc

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
#include "ThreadPool.h"
1919
#include "paddle/fluid/framework/blocking_queue.h"
2020
#include "paddle/fluid/operators/reader/blocking_queue.h"
21+
#include "paddle/fluid/operators/reader/buffered_reader.h"
2122
#include "paddle/fluid/operators/reader/reader_op_registry.h"
2223

2324
namespace paddle {
@@ -232,12 +233,17 @@ class OpenFilesOp : public framework::OperatorBase {
232233
container.reset(new OrderedReaderContainer());
233234
} else {
234235
container.reset(new PreemptiveReaderContainer(
235-
std::min(file_names.size(),
236-
static_cast<size_t>(std::thread::hardware_concurrency()))));
236+
static_cast<size_t>(Attr<int>("thread_num"))));
237237
}
238238

239-
out->Reset(
240-
std::make_shared<MultiFileReader>(file_names, std::move(container)));
239+
std::shared_ptr<framework::ReaderBase> reader(
240+
new MultiFileReader(file_names, std::move(container)));
241+
auto buffer_size = Attr<int>("buffer_size");
242+
if (buffer_size > 1) {
243+
reader = framework::MakeDecoratedReader<BufferedReader>(
244+
reader, platform::CPUPlace(), buffer_size);
245+
}
246+
out->Reset(reader);
241247
}
242248
};
243249

@@ -253,6 +259,8 @@ class OpenFilesOpMaker : public FileReaderMakerBase {
253259
An OpenFilesOp creates a MultiFileReader, which is able to
254260
read data multi-threaded from multiple files.
255261
)DOC");
262+
AddAttr<int>("thread_num", "Number of thread to read files.");
263+
AddAttr<int>("buffer_size", "The reading buffer of these files.");
256264
}
257265
};
258266

0 commit comments

Comments
 (0)