Caffe2 - C++ API
A deep learning, cross platform ML framework
blob_serialization.cc
1 
17 #include "caffe2/core/blob_serialization.h"
18 
19 #include <sstream>
20 #include <mutex>
21 
22 #include "caffe2/core/blob.h"
23 #include "caffe2/utils/proto_utils.h"
24 
25 CAFFE2_DEFINE_int(
26  caffe2_tensor_chunk_size,
27  1000000,
28  "Chunk size to split tensor data into");
29 
30 CAFFE2_DEFINE_int(
31  caffe2_max_tensor_serializer_threads,
32  16,
33  "Maximal number of threads that can be used for tensor serialization");
34 
35 CAFFE2_DEFINE_bool(
36  caffe2_serialize_fp16_as_bytes,
37  false,
38  "Serialize FLOAT16 tensors using byte_data field");
39 
40 namespace caffe2 {
48  public:
49  StringSerializer() {}
50  ~StringSerializer() {}
55  void Serialize(
56  const Blob& blob,
57  const string& name,
58  SerializationAcceptor acceptor) override {
59  CAFFE_ENFORCE(blob.IsType<std::string>());
60 
61  BlobProto blob_proto;
62  blob_proto.set_name(name);
63  blob_proto.set_type("std::string");
64  blob_proto.set_content(blob.template Get<std::string>());
65  acceptor(name, blob_proto.SerializeAsString());
66  }
67 };
68 
74  public:
75  void Deserialize(const BlobProto& proto, Blob* blob) override {
76  *blob->GetMutable<std::string>() = proto.content();
77  }
78 };
79 
80 // The blob serialization member function implementation.
82  const string& name,
83  BlobSerializerBase::SerializationAcceptor acceptor,
84  int chunk_size) const {
85  std::unique_ptr<BlobSerializerBase> serializer(CreateSerializer(meta_.id()));
86  CAFFE_ENFORCE(serializer, "No known serializer for ", meta_.name());
87  serializer->SerializeWithChunkSize(*this, name, acceptor, chunk_size);
88 }
89 
90 // The blob serialization member function implementation.
91 std::string Blob::Serialize(const string& name) const {
92  std::string data;
93  BlobSerializerBase::SerializationAcceptor acceptor = [&data](
94  const std::string&, const std::string& blob) {
95  DCHECK(data.empty()); // should be called once with kNoChunking
96  data = blob;
97  };
98  this->Serialize(name, acceptor, kNoChunking);
99  return data;
100 }
101 
102 // Specialization for StoreDeviceDetail for CPU - nothing needs to be done.
103 template <>
105  const Tensor<CPUContext>& /*input*/,
106  TensorProto* /*proto*/) {}
107 
108 // The actual serialization registry objects.
109 CAFFE_DEFINE_TYPED_REGISTRY(
110  BlobSerializerRegistry,
111  CaffeTypeId,
113  std::unique_ptr);
114 
115 CAFFE_DEFINE_REGISTRY(BlobDeserializerRegistry, BlobDeserializerBase);
116 
117 void Blob::Deserialize(const string& content) {
118  BlobProto blob_proto;
119  CAFFE_ENFORCE(
120  blob_proto.ParseFromString(content),
121  "Cannot parse content into a BlobProto.");
122  Deserialize(blob_proto);
123 }
124 
125 void Blob::Deserialize(const BlobProto& blob_proto) {
126  if (blob_proto.type() == kTensorBlobType) {
127  // This is a tensor object. Depending on the device type, we will
128  // use the corresponding TensorDeserializer.
129  auto deserializer = CreateDeserializer(
130  "Tensor" +
131  DeviceTypeName(blob_proto.tensor().device_detail().device_type()));
132  // Tensor's deserializer should always be registered, but we will double
133  // check if it is not null anyway.
134  CAFFE_ENFORCE(deserializer.get());
135  deserializer->Deserialize(blob_proto, this);
136  } else {
137  auto deserializer = CreateDeserializer(blob_proto.type());
138  CAFFE_ENFORCE(
139  deserializer.get(),
140  "No registered deserializer for type ",
141  blob_proto.type());
142  deserializer->Deserialize(blob_proto, this);
143  }
144 }
145 
146 namespace {
147 // Serialize TensorCPU.
148 REGISTER_BLOB_SERIALIZER(
149  (TypeMeta::Id<TensorCPU>()),
151 REGISTER_BLOB_DESERIALIZER(TensorCPU, TensorDeserializer<CPUContext>);
152 // Serialize std::string
153 REGISTER_BLOB_SERIALIZER((TypeMeta::Id<std::string>()), StringSerializer);
154 REGISTER_BLOB_DESERIALIZER(std::string, StringDeserializer);
155 } // namespace
156 } // namespace caffe2
Blob is a general container that hosts a typed pointer.
Definition: blob.h:41
TensorSerializer is the serializer for Tensors.
BlobDeserializerBase is an abstract class that deserializes a blob from a BlobProto or a TensorProto...
StringSerializer is the serializer for String.
void Serialize(const string &name, BlobSerializerBase::SerializationAcceptor acceptor, int chunk_size=kDefaultChunkSize) const
Serializes the current blob, if possible.
TensorDeserializer is the deserializer for Tensors.
Copyright (c) 2016-present, Facebook, Inc.
T * GetMutable(bool *is_new_object=nullptr)
Gets a mutable pointer to the stored object.
Definition: blob.h:117
void Deserialize(const string &content)
Deserializes from a string containing either BlobProto or TensorProto.
bool IsType() const
Checks if the content stored in the blob is of type T.
Definition: blob.h:74
void Serialize(const Blob &blob, const string &name, SerializationAcceptor acceptor) override
Serializes a Blob.
StringDeserializer is the deserializer for Strings.
BlobSerializerBase is an abstract class that serializes a blob to a string.