提交 0d698a38 编写于 作者: B barrierye

add batch_size check

上级 426d4c3c
...@@ -69,6 +69,14 @@ bool DataFeed::SetFileList(const std::vector<std::string>& files) { ...@@ -69,6 +69,14 @@ bool DataFeed::SetFileList(const std::vector<std::string>& files) {
return true; return true;
} }
void DataFeed::SetBatchSize(int batch_size) {
if (batch_size <= 0) {
LOG(ERROR) << "error: illegal batch size: " << batch_size;
exit(-1);
}
default_batch_size_ = batch_size;
}
bool DataFeed::PickOneFile(std::string& filename) { bool DataFeed::PickOneFile(std::string& filename) {
std::unique_lock<std::mutex> lock(mutex_for_pick_file_); std::unique_lock<std::mutex> lock(mutex_for_pick_file_);
if (file_idx_ == filelist_.size()) { if (file_idx_ == filelist_.size()) {
......
...@@ -164,7 +164,7 @@ class DataFeed { ...@@ -164,7 +164,7 @@ class DataFeed {
virtual bool SetFileList(const std::vector<std::string>& files); virtual bool SetFileList(const std::vector<std::string>& files);
virtual bool Start() = 0; virtual bool Start() = 0;
virtual int Next() = 0; virtual int Next() = 0;
virtual void SetBatchSize(int batch) { default_batch_size_ = batch; } virtual void SetBatchSize(int batch);
virtual int GetBatchSize() { return batch_size_; } virtual int GetBatchSize() { return batch_size_; }
// for subclass with queue // for subclass with queue
virtual void SetQueueSize(int queue_size) { virtual void SetQueueSize(int queue_size) {
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册