Treelite
reference_serializer.cc
Go to the documentation of this file.
1 
8 #include <treelite/tree.h>
9 #include <dmlc/io.h>
10 #include <dmlc/serializer.h>
11 
12 namespace dmlc {
13 namespace serializer {
14 
15 template <typename T>
16 struct Handler<treelite::ContiguousArray<T>> {
17  inline static void Write(Stream* strm, const treelite::ContiguousArray<T>& data) {
18  uint64_t sz = static_cast<uint64_t>(data.Size());
19  strm->Write(sz);
20  strm->Write(data.Data(), sz * sizeof(T));
21  }
22 
23  inline static bool Read(Stream* strm, treelite::ContiguousArray<T>* data) {
24  uint64_t sz;
25  bool status = strm->Read(&sz);
26  if (!status) {
27  return false;
28  }
29  data->Resize(sz);
30  return strm->Read(data->Data(), sz * sizeof(T));
31  }
32 };
33 
34 } // namespace serializer
35 } // namespace dmlc
36 
37 namespace treelite {
38 
39 void Tree::ReferenceSerialize(dmlc::Stream* fo) const {
40  fo->Write(num_nodes);
41  fo->Write(leaf_vector_);
42  fo->Write(leaf_vector_offset_);
43  fo->Write(left_categories_);
44  fo->Write(left_categories_offset_);
45  uint64_t sz = static_cast<uint64_t>(nodes_.Size());
46  fo->Write(sz);
47  fo->Write(nodes_.Data(), sz * sizeof(Tree::Node));
48 
49  // Sanity check
50  CHECK_EQ(nodes_.Size(), num_nodes);
51  CHECK_EQ(nodes_.Size() + 1, leaf_vector_offset_.Size());
52  CHECK_EQ(leaf_vector_offset_.Back(), leaf_vector_.Size());
53  CHECK_EQ(nodes_.Size() + 1, left_categories_offset_.Size());
54  CHECK_EQ(left_categories_offset_.Back(), left_categories_.Size());
55 }
56 
57 void Model::ReferenceSerialize(dmlc::Stream* fo) const {
58  fo->Write(num_feature);
59  fo->Write(num_output_group);
60  fo->Write(random_forest_flag);
61  fo->Write(&param, sizeof(param));
62  uint64_t sz = static_cast<uint64_t>(trees.size());
63  fo->Write(sz);
64  for (const Tree& tree : trees) {
65  tree.ReferenceSerialize(fo);
66  }
67 }
68 
69 } // namespace treelite
model structure for tree ensemble
Definition: tree.h:27