caffe源码分析-DataReader

bear_fish

发布于 2019-02-25 11:53:00

4740

发布于 2019-02-25 11:53:00

DataReader作为DataLayer的数据成员变量，以多线程的方式从数据库(如lmdb, hdf5)读取数据：

class DataReader {
public:
    explicit DataReader(const LayerParameter& param);
    ~DataReader();

    inline BlockingQueue<Datum*>& free() const {
        return queue_pair_->free_;
    }
    inline BlockingQueue<Datum*>& full() const {
        return queue_pair_->full_;
    }

protected:
    // Queue pairs are shared between a body and its readers
    class QueuePair {
    public:
        explicit QueuePair(int size);
        ~QueuePair();

        BlockingQueue<Datum*> free_;
        BlockingQueue<Datum*> full_;

    DISABLE_COPY_AND_ASSIGN(QueuePair);
    };

    // A single body is created per source
    class Body : public InternalThread {
    public:
        explicit Body(const LayerParameter& param);
        virtual ~Body();

    protected:
        void InternalThreadEntry();
        void read_one(db::Cursor* cursor, QueuePair* qp);

        const LayerParameter param_;
        BlockingQueue<shared_ptr<QueuePair> > new_queue_pairs_;

        friend class DataReader;

    DISABLE_COPY_AND_ASSIGN(Body);
    };

    // A source is uniquely identified by its layer name + path, in case
    // the same database is read from two different locations in the net.
    static inline string source_key(const LayerParameter& param) {
        return param.name() + ":" + param.data_param().source();
    }

protected:
    const shared_ptr<QueuePair> queue_pair_;
    shared_ptr<Body> body_;

    static map<const string, boost::weak_ptr<DataReader::Body> > bodies_;

    DISABLE_COPY_AND_ASSIGN(DataReader);
};

下面看构造函数，预读取的总的数据长度为prefetch*batch_size， weak_ptr结合shared_ptr保证线程安全：

DataReader::DataReader(const LayerParameter& param)
        : queue_pair_(new QueuePair(  //
        param.data_param().prefetch() * param.data_param().batch_size())) {
    // Get or create a body
    boost::mutex::scoped_lock lock(bodies_mutex_);
    string key = source_key(param);
    weak_ptr<Body>& weak = bodies_[key];
    body_ = weak.lock();
    if (!body_) {
        body_.reset(new Body(param));
        bodies_[key] = weak_ptr<Body>(body_);
    }
    body_->new_queue_pairs_.push(queue_pair_);
}

线程的核心函数：

void DataReader::Body::InternalThreadEntry() {
    shared_ptr<db::DB> db(db::GetDB(param_.data_param().backend()));
    db->Open(param_.data_param().source(), db::READ);
    shared_ptr<db::Cursor> cursor(db->NewCursor());
    vector<shared_ptr<QueuePair> > qps;

        // Main loop
        while (!must_stop()) {
            for (int i = 0; i < solver_count; ++i) {
                read_one(cursor.get(), qps[i].get());
            }
            CHECK_EQ(new_queue_pairs_.size(), 0);
        }
}

void DataReader::Body::read_one(db::Cursor* cursor, QueuePair* qp) {
    Datum* datum = qp->free_.pop();
    datum->ParseFromString(cursor->value());
    qp->full_.push(datum);
    // go to the next iter
    cursor->Next();
}