clstm-hdf5-to-proto.cc

// Loading and saving networks with HDF5. This is obsolete, use the
// protobuf based implementation instead.

#include "clstm.h"
#include "h5eigen.h"
#include <assert.h>
#include <iostream>
#include <vector>
#include <string>
#include <memory>
#include <math.h>
#include <Eigen/Dense>
#include "clstm.pb.h"
#include <stdarg.h>

namespace {
inline void print() { std::cout << std::endl; }

template <class T>
inline void print(const T &arg) {
  std::cout << arg << std::endl;
}

template <class T, typename... Args>
inline void print(T arg, Args... args) {
  std::cout << arg << " ";
  print(args...);
}
}

namespace ocropus {
map<string, INetworkFactory> network_factories;

int register_stacked(const string &name, function<void(INetwork *)> f) {
  network_factories[name] = [f]() {
    INetwork *stacked = make_Stacked();
    stacked->initializer = f;
    return stacked;
  };
  return 0;
}

#define C(X, Y) X##Y
#define REGISTER_STACKED(X) \
  int C(status_, X) = register_stacked(#X, C(init_, X));

void init_LSTM1(INetwork *net) {
  net->name = "lstm1";
  net->attributes["kind"] = "lstm1";
  int ni = net->irequire("ninput");
  int nh = net->irequire("nhidden");
  int no = net->irequire("noutput");
  Network fwd, softmax;
  fwd.reset(make_LSTM());
  fwd->init(nh, ni);
  net->add(fwd);
  softmax.reset(make_SoftmaxLayer());
  softmax->init(no, nh);
  net->add(softmax);
}
REGISTER_STACKED(LSTM1);

void init_REVLSTM1(INetwork *net) {
  net->name = "revlstm1";
  int ni = net->irequire("ninput");
  int nh = net->irequire("nhidden");
  int no = net->irequire("noutput");
  Network fwd, rev, softmax;
  fwd.reset(make_LSTM());
  fwd->init(nh, ni);
  rev.reset(make_Reversed());
  rev->add(fwd);
  net->add(rev);
  softmax.reset(make_SoftmaxLayer());
  softmax->init(no, nh);
  net->add(softmax);
}
REGISTER_STACKED(REVLSTM1);

#if 0
void init_BIDILSTM(INetwork *net) {
    net->name = "bidilstm";
    net->attributes["kind"] = "bidi";
    int ni = net->irequire("ninput");
    int nh = net->irequire("nhidden");
    int no = net->irequire("noutput");
    Network fwd, bwd, parallel, reversed, softmax;
    fwd = make_net(net->attr("lstm_type", "LSTM"));
    bwd = make_net(net->attr("lstm_type", "LSTM"));
    softmax = make_net(net->attr("output_type", "SoftmaxLayer"));
    fwd->init(nh, ni);
    bwd->init(nh, ni);
    reversed.reset(make_Reversed());
    reversed->add(bwd);
    parallel.reset(make_Parallel());
    parallel->add(fwd);
    parallel->add(reversed);
    net->add(parallel);
    softmax->init(no, 2*nh);
    net->add(softmax);
}
REGISTER_STACKED(BIDILSTM);

void init_LRBIDILSTM(INetwork *net) {
    net->set("output_type", "SigmoidLayer");
    init_BIDILSTM(net);
}
REGISTER_STACKED(LRBIDILSTM);

void init_LINBIDILSTM(INetwork *net) {
    net->set("output_type", "LinearLayer");
    init_BIDILSTM(net);
}
REGISTER_STACKED(LINBIDILSTM);

void init_BIDILSTM2(INetwork *net) {
    net->name = "bidilstm2";
    net->attributes["kind"] = "bidi2";
    int ni = net->irequire("ninput");
    int nh = net->irequire("nhidden");
    int nh2 = net->irequire("nhidden2");
    int no = net->irequire("noutput");
    // cerr << ">>> " << ni << " " << nh << " " << nh2 << " " << no << endl;
    Network parallel1, parallel2, logreg;
    parallel1.reset(make_BidiLayer());
    parallel1->init(nh, ni);
    net->add(parallel1);
    parallel2.reset(make_BidiLayer());
    parallel2->init(nh2, 2*nh);
    net->add(parallel2);
    logreg.reset(make_SoftmaxLayer());
    logreg->init(no, 2*nh2);
    net->add(logreg);
};
REGISTER_STACKED(BIDILSTM2);
#endif

static string fixup_net_name(string kind) {
  if (kind == "bidi") return "BIDILSTM";
  if (kind == "lrbidi") return "LRBIDILSTM";
  if (kind == "bidi2") return "BIDILSTM2";
  if (kind == "lstm1") return "LSTM1";
  if (kind == "uni") return "LSTM1";
  return kind;
}

Network make_net(string kind) {
  if (kind == "") THROW("empty network kind in make_net");
  Network net;
  net = make_layer(kind);
  if (net) return net;
  kind = fixup_net_name(kind);
  auto it = network_factories.find(kind);
  if (it == network_factories.end()) THROW("unknown network kind");
  net.reset(it->second());
  net->attributes["kind"] = kind;
  return net;
}
}

namespace ocropus {
using std::cout;
using std::endl;

void save_codec(h5eigen::HDF5 *h5, const char *name, vector<int> &codec,
                int n) {
  assert(codec.size() == 0 || codec.size() == n);
  Vec vcodec;
  vcodec.resize(n);
  if (codec.size() == 0)
    for (int i = 0; i < n; i++) vcodec[i] = i;
  else
    for (int i = 0; i < n; i++) vcodec[i] = codec[i];

  h5->put(vcodec, name);
}

void load_codec(vector<int> &codec, h5eigen::HDF5 *h5, const char *name) {
  codec.resize(0);
  if (!h5->exists(name)) return;
  Vec vcodec;
  h5->get(vcodec, name);
  codec.resize(vcodec.size());
  for (int i = 0; i < vcodec.size(); i++) codec[i] = vcodec[i];
}

void save_net_raw(const char *fname, INetwork *net) {
  using namespace h5eigen;
  unique_ptr<HDF5> h5(make_HDF5());
  h5->open(fname, true);
  net->attributes["clstm-version"] = "1";
  for (auto &kv : net->attributes) {
    h5->setAttr(kv.first, kv.second);
  }
  save_codec(h5.get(), "codec", net->codec, net->noutput());
  save_codec(h5.get(), "icodec", net->icodec, net->ninput());

  net->weights("", [&h5](const string &prefix, VecMat a, VecMat da) {
    if (a.mat)
      h5->put(*a.mat, prefix.c_str());
    else if (a.vec)
      h5->put(*a.vec, prefix.c_str());
    else
      THROW("oops (save type)");
  });
}

void load_net_raw(INetwork *net, const char *fname) {
  using namespace h5eigen;
  unique_ptr<HDF5> h5(make_HDF5());
  h5->open(fname);
  h5->getAttrs(net->attributes);
  load_codec(net->icodec, h5.get(), "icodec");
  load_codec(net->codec, h5.get(), "codec");
  net->weights("", [&h5](const string &prefix, VecMat a, VecMat da) {
    if (a.mat)
      h5->get(*a.mat, prefix.c_str());
    else if (a.vec)
      h5->get1d(*a.vec, prefix.c_str());
    else
      THROW("oops (load type)");
  });
  net->networks("", [](string s, INetwork *net) { net->postLoad(); });
}

void load_attributes(map<string, string> &attributes, const string &fname) {
  using namespace h5eigen;
  unique_ptr<HDF5> h5(make_HDF5());
  h5->open(fname.c_str());
  h5->getAttrs(attributes);
}

Network load_net_hdf5(const string &fname) {
  bool verbose = (getenv("load_verbose") && atoi(getenv("load_verbose")));
  map<string, string> attributes;
  load_attributes(attributes, fname);
  string kind = attributes["kind"];
  if (verbose) {
    cout << "loading " << fname << endl;
    cout << endl;
    cout << "attributes:" << endl;
    for (auto kv : attributes)
      cout << "   " << kv.first << " = " << kv.second << endl;
    cout << endl;
  }
  Network net;
  net = make_net(kind);
  net->attributes = attributes;
  net->initialize();
  load_net_raw(net.get(), fname.c_str());
  if (verbose) {
    cout << "network:" << endl;
    net->info("");
    cout << endl;
  }
  return net;
}

void save_net_hdf5(const string &fname, Network net) {
  rename(fname.c_str(), (fname + "~").c_str());
  save_net_raw(fname.c_str(), net.get());
}
}

using namespace ocropus;
using namespace std;

int main(int argc, char **argv) {
  Network net = load_net_hdf5(argv[1]);
  save_as_proto(argv[2], net.get());
}