#ifndef CAFFE2_CORE_BLOB_SERIALIZATION_H_ #define CAFFE2_CORE_BLOB_SERIALIZATION_H_ #include #include #include #include "caffe2/core/blob.h" #include "caffe2/core/blob_serializer_base.h" #include "caffe2/core/tensor.h" #include #include "caffe2/core/types.h" #include "caffe2/utils/simple_queue.h" C10_DECLARE_int(caffe2_tensor_chunk_size); C10_DECLARE_int(caffe2_max_tensor_serializer_threads); C10_DECLARE_bool(caffe2_serialize_fp16_as_bytes); namespace caffe2 { constexpr auto kTensorBlobType = "Tensor"; // String used to separate chunk id from the blob name when storing in DB constexpr auto kChunkIdSeparator = "#%"; /** * Serializes the given blob, if possible. Note that this serialization uses * the registration mechanism and one has to implement specific serialization * approaches for specific classes. Acceptor should take care of writing data * to the actual storage. */ CAFFE2_API void SerializeBlob( const Blob& blob, const string& name, BlobSerializerBase::SerializationAcceptor acceptor, int chunk_size = kDefaultChunkSize); /** * @brief Convenience function to serialize a blob to a string. * * This is a conveinence function to serialize small Blobs that produce * manageable serialized strings. To serialize big blobs such as * large sparse tensors, use the fully-functional interface in * blob_serializer_base.h. * * NOTE: this function doesn't do chunking and might break with big tensors. */ CAFFE2_API string SerializeBlob(const Blob& blob, const string& name); /** * Deserializes from a string containing either BlobProto or TensorProto. If * the deserialization fails, the content in the blob should no longer be * trusted. */ CAFFE2_API void DeserializeBlob(const string& content, Blob* result); CAFFE2_API void DeserializeBlob(const BlobProto& proto, Blob* result); /* * Get an empty Tensor from the TensorProto given the meta data in proto (data * type and size of the Tensor) without actually filling in the data. * * We need this function because we want to construct a fully initialized Tensor * in the beginning instead of keeping partially initialized Tensor around the * process. Consider the case when we have a Tensor that is split into multiple * protos during serialization, in deserialization, we have to fill the Tensor * in multiple calls to Deserialize, therefore we need to create a new Tensor * with the correct size and data type before the call to Deserialize, because * otherwise we will have to check whether the function call is the first call * to initialize the underlying Tensor, which makes the function stateful and * complicated. * * The legacy code get away with this problem by passing in a partially * initialized Tensor and use Resize and mutable_data to set the correct size, * data type and allocate memory for the Tensor, so the state is encoded in * these function calls. e.g. mutable_data will allocate memory on the first * call and it will return a pointer to the allocated memory on later calls. */ CAFFE2_API Tensor EmptyTensorFromProto(const TensorProto& proto); /** * @brief TensorSerializer is the serializer for Tensors. * * TensorSerializer takes in a blob that contains a Tensor, and serializes it * into a TensorProto protocol buffer. */ class CAFFE2_API TensorSerializer : public BlobSerializerBase { public: TensorSerializer() {} ~TensorSerializer() override {} /** * Serializes a Blob. Note that this blob has to contain Tensor, * otherwise this function produces a fatal error. */ void Serialize( const void* pointer, TypeMeta typeMeta, const string& name, SerializationAcceptor acceptor) override; void SerializeWithChunkSize( const void* pointer, TypeMeta typeMeta, const string& name, SerializationAcceptor acceptor, int chunk_size) override; void Serialize( const Tensor& tensor, const string& name, TensorProto* proto, size_t chunkBegin, int32_t chunkSize); private: // A utility function to store the device context detauls. void StoreDeviceDetail(const Tensor& input, TensorProto* proto); unique_ptr context_; }; /** * @brief TensorDeserializer is the deserializer for Tensors. * * The device that the deserialized Tensor will live under is determined by the * device_detail field. If you want to specify the device of the deserialized * tensor, change the TensorProto's corresponding fields before calling * Deserialize. */ class CAFFE2_API TensorDeserializer : public BlobDeserializerBase { public: void Deserialize(const BlobProto& proto, Blob* blob) override; /* There are cases when a Tensor is split into multiple protos and * we have to call Deserialize multiple times to get the complete deserialized * Tensor, each call will fill part of the Tensor given the segment begin and * end information in proto, therefore we have to pass in the Tensor pointer * rather than create a new Tensor everytime. * * Precondition: Tensor must be initialized */ void DeserializeToTensor(const TensorProto& proto, Tensor* tensor); /* Deserialize the proto and return a new Tensor * This is a utility function that combines EmptyTensorFromProto and * Deserialize(const TensorProto&, Tensor*); */ Tensor Deserialize(const TensorProto& proto); }; //////////////////////////////////////////////////////////////////////////////// // Implementations //////////////////////////////////////////////////////////////////////////////// namespace detail { template inline void CopyToProtoAsIs( const size_t size, const SrcType* src, google::protobuf::RepeatedField* field, BaseContext* context) { static_assert( sizeof(SrcType) == sizeof(DstType), "The source type and dest type cannot be copied as-is. Did " "you mean CopyToProtoWithCast?"); field->Reserve(size); for (size_t i = 0; i < size; ++i) { field->Add(0); } context->template CopyToCPU( size, src, reinterpret_cast(field->mutable_data())); // Make sure that we finish the copy into the protobuf. context->FinishDeviceComputation(); } template inline void CopyToProtoWithCast( const size_t size, const SrcType* src, google::protobuf::RepeatedField* field, BaseContext* context) { // TODO: we are having one unnecessary copy here if the context is already // CPUContext. Remove it if it is performance critical. unique_ptr buffer(new SrcType[size]); context->template CopyToCPU(size, src, buffer.get()); context->FinishDeviceComputation(); field->Reserve(size); for (size_t i = 0; i < size; ++i) { field->Add(static_cast(buffer[i])); } } template inline void CopyFromProtoAsIs( const size_t size, const google::protobuf::RepeatedField& field, DstType* dst, BaseContext* context) { static_assert( sizeof(SrcType) == sizeof(DstType), "The source type and dest type cannot be copied as-is. Did " "you mean CopyFromProtoWithCast?"); CAFFE_ENFORCE_EQ(size, field.size(), "Incorrect proto field size."); context->template CopyFromCPU( size, reinterpret_cast(field.data()), dst); } template inline void CopyFromProtoWithCast( const size_t size, const google::protobuf::RepeatedField& field, DstType* dst, BaseContext* context) { CAFFE_ENFORCE_EQ(size, field.size(), "Incorrect proto field size."); // TODO: we are having one unnecessary copy here if the context is already // CPUContext. Remove it if it is performance critical. unique_ptr buffer(new DstType[size]); const SrcType* src = field.data(); for (size_t i = 0; i < size; ++i) { buffer[i] = static_cast(src[i]); } context->template CopyFromCPU(size, buffer.get(), dst); } } // namespace detail //////////////////////////////////////////////////////////////////////////////// // Serialization Helpers //////////////////////////////////////////////////////////////////////////////// // Converts MessageLite to string while also checking that SerializeAsString // succeeds. Pass description of class/function of the call if you'd // like it appended to the error message. CAFFE2_API std::string SerializeAsString_EnforceCheck( const google::protobuf::MessageLite&, const char* error_location = nullptr); // Convert BlobProto to string with success checks. inline std::string SerializeBlobProtoAsString_EnforceCheck( const BlobProto& blob) { return SerializeAsString_EnforceCheck(blob, blob.name().c_str()); } } // namespace caffe2 #endif // CAFFE2_CORE_BLOB_SERIALIZATION_H_