blob_serializer_base.h

#pragma once

#include <string>
#include <functional>

#include "c10/util/Registry.h"
#include "caffe2/core/common.h"
#include "caffe2/proto/caffe2_pb.h"

namespace caffe2 {

class Blob;

constexpr int kDefaultChunkSize = -1;
constexpr int kNoChunking = 0;

class BlobSerializerBase {
 public:
  virtual ~BlobSerializerBase() {}
  using SerializationAcceptor =
     std::function<void(const std::string& blobName, const std::string& data)>;
  virtual void Serialize(
      const void* pointer,
      TypeMeta typeMeta,
      const std::string& name,
      SerializationAcceptor acceptor) = 0;

  virtual void SerializeWithChunkSize(
      const void* pointer,
      TypeMeta typeMeta,
      const std::string& name,
      SerializationAcceptor acceptor,
      int /*chunk_size*/) {
    // Base implementation.
    Serialize(pointer, typeMeta, name, acceptor);
  }
};

// The Blob serialization registry and serializer creator functions.
C10_DECLARE_TYPED_REGISTRY(
    BlobSerializerRegistry,
    TypeIdentifier,
    BlobSerializerBase,
    std::unique_ptr);
#define REGISTER_BLOB_SERIALIZER(id, ...) \
  C10_REGISTER_TYPED_CLASS(BlobSerializerRegistry, id, __VA_ARGS__)
// Creates an operator with the given operator definition.
inline unique_ptr<BlobSerializerBase> CreateSerializer(TypeIdentifier id) {
  return BlobSerializerRegistry()->Create(id);
}


class CAFFE2_API BlobDeserializerBase {
 public:
  virtual ~BlobDeserializerBase() {}

  // Deserializes from a BlobProto object.
  virtual void Deserialize(const BlobProto& proto, Blob* blob) = 0;
};

C10_DECLARE_REGISTRY(BlobDeserializerRegistry, BlobDeserializerBase);
#define REGISTER_BLOB_DESERIALIZER(name, ...) \
  C10_REGISTER_CLASS(BlobDeserializerRegistry, name, __VA_ARGS__)
// Creates an operator with the given operator definition.
inline unique_ptr<BlobDeserializerBase> CreateDeserializer(const string& type) {
  return BlobDeserializerRegistry()->Create(type);
}


} // namespace caffe2