caffe源码分析-db, io

bear_fish

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

5220

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

本文主要分析下caffe的源码，io操作，以及数据库文件（如lmdb）读取。

例如，从prototxt读取网络初始化参数：

NetParameter param;
ReadNetParamsFromTextFileOrDie(path, &param);

bool ReadProtoFromTextFile(const char *filename, Message *proto) {
    int fd = open(filename, O_RDONLY);
    CHECK_NE(fd, -1) << "File not found: " << filename;
    FileInputStream *input = new FileInputStream(fd);
    bool success = google::protobuf::TextFormat::Parse(input, proto);
    delete input;
    close(fd);
    return success;
}

例如读取下面这个net

name: "LeNet"
layer {
  name: "mnist"
  type: "Data"
  top: "data"
  top: "label"
  include {
    phase: TRAIN
  }
  transform_param {
    scale: 0.00390625
  }
  data_param {
    source: "/home/xy/caffe-master/examples/mnist/mnist_train_lmdb"
    batch_size: 64
    backend: LMDB
  }
}

将参数写为prototxt文件：

void WriteProtoToTextFile(const Message &proto, const char *filename) {
    int fd = open(filename, O_WRONLY | O_CREAT | O_TRUNC, 0644);
    FileOutputStream *output = new FileOutputStream(fd);
    CHECK(google::protobuf::TextFormat::Print(proto, output));
    delete output;
    close(fd);
}

输出网络(打印), 示例:

// read proto txt to proto class
NetParameter param;
ReadNetParamsFromTextFileOrDie(path, &param);

// print proto class
PrintProto(param);

函数定义如下：

void PrintProto(const Message &proto) {
    OstreamOutputStream os(&std::cout);
    google::protobuf::TextFormat::Print(proto, &os);
}

读取proto文件转化为Datum：

// read Datum from file
string datum_path = base_dir + "0.proto";
Datum datum;
ReadProtoFromBinaryFileOrDie(datum_path, &datum);

bool ReadProtoFromBinaryFile(const char *filename, Message *proto) {
    int fd = open(filename, O_RDONLY);
    CHECK_NE(fd, -1) << "File not found: " << filename;
    ZeroCopyInputStream *raw_input = new FileInputStream(fd);
    CodedInputStream *coded_input = new CodedInputStream(raw_input);
    coded_input->SetTotalBytesLimit(kProtoReadBytesLimit, 536870912);

    bool success = proto->ParseFromCodedStream(coded_input);

    delete coded_input;
    delete raw_input;
    close(fd);
    return success;
}

Datum与cv::Mat的转化:

cv::Mat DecodeDatumToCVMatNative(const Datum &datum) {
    cv::Mat cv_img;
    CHECK(datum.encoded()) << "Datum not encoded";
    const string &data = datum.data();
    std::vector<char> vec_data(data.c_str(), data.c_str() + data.size());
    cv_img = cv::imdecode(vec_data, -1);
    if (!cv_img.data) {
        LOG(ERROR) << "Could not decode datum ";
    }
    return cv_img;
}

void CVMatToDatum(const cv::Mat &cv_img, Datum *datum) {
    datum->set_channels(cv_img.channels());
    datum->set_height(cv_img.rows);
    datum->set_width(cv_img.cols);
    datum->clear_data();
    datum->clear_float_data();
    datum->set_encoded(false);
    int datum_channels = datum->channels();
    int datum_height = datum->height();
    int datum_width = datum->width();
    int datum_size = datum_channels * datum_height * datum_width;
    std::string buffer(datum_size, ' ');
    for (int h = 0; h < datum_height; ++h) {
        const uchar *ptr = cv_img.ptr<uchar>(h);
        int img_index = 0;
        for (int w = 0; w < datum_width; ++w) {
            for (int c = 0; c < datum_channels; ++c) {
                int datum_index = (c * datum_height + h) * datum_width + w;
                buffer[datum_index] = static_cast<char>(ptr[img_index++]);
            }
        }
    }
    datum->set_data(buffer);
}

下面简要看db：

class Cursor {
public:
    Cursor() { }
    virtual ~Cursor() { }
    virtual void SeekToFirst() = 0;
    virtual void Next() = 0;
    virtual string key() = 0;
    virtual string value() = 0;
    virtual bool valid() = 0;

    DISABLE_COPY_AND_ASSIGN(Cursor);
};

    class Transaction {
    public:
        Transaction() { }
        virtual ~Transaction() { }
        virtual void Put(const string& key, const string& value) = 0;
        virtual void Commit() = 0;

    DISABLE_COPY_AND_ASSIGN(Transaction);
    };

    class DB {
    public:
        DB() { }
        virtual ~DB() { }
        virtual void Open(const string& source, Mode mode) = 0;
        virtual void Close() = 0;
        virtual Cursor* NewCursor() = 0;
        virtual Transaction* NewTransaction() = 0;

    DISABLE_COPY_AND_ASSIGN(DB);
    };

    DB* GetDB(DataParameter::DB backend);
    DB* GetDB(const string& backend);
}

DB *GetDB(const string &backend) {
    if (backend == "lmdb") {
        return new LMDB();
    }
    LOG(FATAL) << "Unknown database backend";
    return NULL;
}