Serialization in C++ Frontend API

kolkir: <p>Hello,<br>
I have a question about using serialization API in C++. I want to load model’s weights in C++ program, and save them in python one.</p>
<p>In python I’m doing next:</p>
<pre><code class="lang-auto">class MaskRCNN(nn.Module):
...
  def save_weights():
    torch.save(self.state_dict(), self.checkpoint_path.format(epoch))
...
</code></pre>
<p>In C++ I’m trying to load saved weights in the next way:</p>
<pre><code class="lang-auto">torch::serialize::InputArchive archive;
archive.load_from(params_path);
</code></pre>
<p>But this code fails with exception:</p>
<pre><code class="lang-auto">tag == RecordTags::FOOTER ASSERT FAILED at /development/lib/sources/pytorch/caffe2/serialize/inline_container.h:234, please report a bug to PyTorch. File footer has wrong record type. Is this file corrupted? (readAndValidateFileFooter at /development/lib/sources/pytorch/caffe2/serialize/inline_container.h:234)
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt; const&amp;) + 0x6c (0x7f6a66008e3c in /home/kirill/development/lib/lib/libc10.so)
frame #1: &lt;unknown function&gt; + 0x5ccc90 (0x7f6a7ce8dc90 in /home/kirill/development/lib/lib/libtorch.so.1)
frame #2: &lt;unknown function&gt; + 0x5cd6b4 (0x7f6a7ce8e6b4 in /home/kirill/development/lib/lib/libtorch.so.1)
frame #3: torch::jit::load(std::istream&amp;) + 0x2af (0x7f6a7ce8c6ff in /home/kirill/development/lib/lib/libtorch.so.1)
frame #4: torch::jit::load(std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt; const&amp;) + 0x51 (0x7f6a7ce8c961 in /home/kirill/development/lib/lib/libtorch.so.1)
frame #5: torch::serialize::InputArchive::load_from(std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt; const&amp;) + 0x21 (0x7f6a7d057c91 in /home/kirill/development/lib/lib/libtorch.so.1)
frame #6: main + 0x5ce (0x55bd69bd3538 in /home/kirill/development/mlcpp/mask_rcnn_pytorch/build/mask-rcnn)
frame #7: __libc_start_main + 0xf3 (0x7f6a64c7d223 in /usr/lib/libc.so.6)
frame #8: _start + 0x2e (0x55bd69bd2bae in /home/kirill/development/mlcpp/mask_rcnn_pytorch/build/mask-rcnn)
</code></pre>
<p>In python and can successfully load and use saved weights in the next way</p>
<pre><code class="lang-auto">model.load_state_dict(torch.load(weights_file_name))
</code></pre>
<p>Could you tell me if I misuse API, and it’s incompatible? What functions I should use to solve my problem, if they a exist now?</p>
<p>Thanks, in advance.</p>

kolkir: <p>After some investigation I came up with a next solution: create a dictionary of lists with tensor values extracted from  models state_dict, and save it as json, then parse this file in C++ and initialize <code>torch::Tensor</code>  objects.  Approach with <code>torch.load\save</code> doesn’t work in C++ because these functions are based on Python pickle. There are two other approaches to share models data between languages: <code>torch.jit</code> and <code>onnx</code>. But they relies on model execution tracing and it’s not easy to use them in case complex models like Mask-RCNN for example.</p>
<p>Implementation:</p>
<ol>
<li>In python:</li>
</ol>
<pre><code class="lang-auto"> raw_state_dict = {}
 for k, v in model.state_dict().items():
     if isinstance(v, torch.Tensor):
         raw_state_dict[k] = (list(v.size()), v.numpy().tolist())
         break
     else:
         print("State parameter type error : {}".format(k))
         exit(-1)

 with open('mask_rcnn_coco.json', 'w') as outfile:
     json.dump(raw_state_dict, outfile)
</code></pre>
<ol start="2">
<li>In C++:</li>
</ol>
<pre><code class="lang-auto">
#include &lt;rapidjson/error/en.h&gt;
#include &lt;rapidjson/filereadstream.h&gt;
#include &lt;rapidjson/reader.h&gt;

#include &lt;iostream&gt;
#include &lt;stack&gt;

namespace {

enum class ReadState {
  None,
  DictObject,
  ParamName,
  SizeTensorPair,
  TensorSize,
  SizeTensorPairDelim,
  TensorValue,
  List
};

struct DictHandler
    : public rapidjson::BaseReaderHandler&lt;rapidjson::UTF8&lt;&gt;, DictHandler&gt; {
  DictHandler() {}

  bool Double(double d) {
    std::cout &lt;&lt; "double " &lt;&lt; d &lt;&lt; std::endl;
    if (current_state_.top() == ReadState::List ||
        current_state_.top() == ReadState::TensorValue) {
      blob_.push_back(static_cast&lt;float&gt;(d));
      ++index_;
    } else {
      throw std::logic_error("Double parsing error");
    }
    return true;
  }

  bool Uint(unsigned u) {
    std::cout &lt;&lt; "uint " &lt;&lt; u &lt;&lt; std::endl;
    if (current_state_.top() == ReadState::List ||
        current_state_.top() == ReadState::TensorValue) {
      blob_.push_back(static_cast&lt;float&gt;(u));
      ++index_;
    } else if (current_state_.top() == ReadState::TensorSize) {
      size_.push_back(static_cast&lt;int64_t&gt;(u));
    } else {
      throw std::logic_error("UInt parsing error");
    }
    return true;
  }

  bool Key(const char* str, rapidjson::SizeType length, bool /*copy*/) {
    key_.assign(str, length);
    std::cout &lt;&lt; "key " &lt;&lt; key_ &lt;&lt; std::endl;
    if (current_state_.top() == ReadState::DictObject) {
      current_state_.push(ReadState::ParamName);
    } else {
      throw std::logic_error("Key parsing error");
    }
    return true;
  }

  bool StartObject() {
    std::cout &lt;&lt; "start object" &lt;&lt; std::endl;
    if (current_state_.top() == ReadState::None) {
      current_state_.pop();
      current_state_.push(ReadState::DictObject);
    } else {
      throw std::logic_error("Start object parsing error");
    }
    return true;
  }

  bool EndObject(rapidjson::SizeType /*memberCount*/) {
    std::cout &lt;&lt; "end object" &lt;&lt; std::endl;
    if (current_state_.top() != ReadState::DictObject) {
      throw std::logic_error("End object parsing error");
    }
    return true;
  }

  void StartData() {
    current_state_.push(ReadState::TensorValue);
    auto total_length = std::accumulate(size_.begin(), size_.end(), 1,
                                        std::multiplies&lt;int64_t&gt;());
    blob_.resize(static_cast&lt;size_t&gt;(total_length));
    blob_.clear();
    index_ = 0;
  }

  bool StartArray() {
    std::cout &lt;&lt; "start array" &lt;&lt; std::endl;
    if (current_state_.top() == ReadState::List) {
      current_state_.push(ReadState::List);
    } else if (current_state_.top() == ReadState::ParamName) {
      current_state_.push(ReadState::SizeTensorPair);
    } else if (current_state_.top() == ReadState::SizeTensorPair) {
      current_state_.push(ReadState::TensorSize);
      size_.clear();
    } else if (current_state_.top() == ReadState::SizeTensorPairDelim) {
      current_state_.pop();
      StartData();
    } else if (current_state_.top() == ReadState::TensorValue) {
      current_state_.push(ReadState::List);
    } else {
      throw std::logic_error("Start array parsing error");
    }
    return true;
  }

  bool EndArray(rapidjson::SizeType elementCount) {
    std::cout &lt;&lt; "end array" &lt;&lt; std::endl;
    if (current_state_.top() == ReadState::List) {
      current_state_.pop();
    } else if (current_state_.top() == ReadState::SizeTensorPair) {
      current_state_.pop();
      assert(current_state_.top() == ReadState::ParamName);
      current_state_.pop();
      std::cout &lt;&lt; "Add new param" &lt;&lt; std::endl;
      dict.push_back({key_, tensor_});
    } else if (current_state_.top() == ReadState::TensorSize) {
      current_state_.pop();
      if (elementCount == 0) {
        size_.push_back(1);
        StartData();
      } else {
        current_state_.push(ReadState::SizeTensorPairDelim);
      }
    } else if (current_state_.top() == ReadState::TensorValue) {
      current_state_.pop();
      assert(index_ == static_cast&lt;int64_t&gt;(blob_.size()));
      at::Tensor tensor_image = torch::from_blob(
          blob_.data(), at::IntList(size_), at::CPU(at::kFloat));
      if (blob_.size() == 1) {
        assert(current_state_.top() == ReadState::SizeTensorPair);
        current_state_.pop();
        assert(current_state_.top() == ReadState::ParamName);
        current_state_.pop();
        std::cout &lt;&lt; "Add new param" &lt;&lt; std::endl;
        dict.push_back({key_, tensor_});
      }
    } else {
      throw std::logic_error("End array parsing error");
    }
    return true;
  }

  std::string key_;
  std::vector&lt;int64_t&gt; size_;
  torch::Tensor tensor_;
  std::vector&lt;float&gt; blob_;
  int64_t index_{0};

  std::stack&lt;ReadState&gt; current_state_{{ReadState::None}};

  std::vector&lt;std::pair&lt;std::string, torch::Tensor&gt;&gt; dict;
};
}  // namespace

std::vector&lt;std::pair&lt;std::string, torch::Tensor&gt;&gt; LoadStateDict(
    const std::string&amp; file_name) {
  auto* file = std::fopen(file_name.c_str(), "r");
  if (file) {
    char readBuffer[65536];
    rapidjson::FileReadStream is(file, readBuffer, sizeof(readBuffer));
    rapidjson::Reader reader;
    DictHandler handler;
    auto res = reader.Parse(is, handler);
    std::fclose(file);

    if (!res) {
      throw std::runtime_error(rapidjson::GetParseError_En(res.Code()));
    }

    return handler.dict;
  }
  return {};
}
</code></pre>

Tushar_Sangam: <p>Hi, I was inspired by your solution, but I found JSON deserializing to be slow with large tensors.<br>
So I used the similar logic but implemented a simple binary read write mechanism,<br>
torch tensors can convert to  numpy and numpy can convert to bytestring that’s how the python can write the underlying allocations data in binary file,<br>
While in C++ the ifstream’s read method in binary mode allows us to pass a char* to char array of size length and reads length bytes in that array.<br>
Since the tensor created in libtorch has same shape and size that of python tensor, its .numel() * .elemenetsize() yeilds the exact number of bytes that needs to be read from the file, and the tensor.sizes() can be passed to stride the inner data accordingly.</p>
<p>If you want I can post the snippet of my loading and saving code.<br>
I think the load times and save times were faster</p>