介绍
Caffe2中Blob的概念应该来自于Caffe。它是有类型的内存抽象,主要包含两个成员,一为指向存储元素的指针,另一则为此元素的类型(TypeMeta)。这么说来它其实与Tensor好像,本质上它有些赘余,更像是来自Caffe的一种包袱。在笔者已知的框架设计里像Tensorflow/Pytorch/Mxnet等无不是只提供Tensor这么一种有类型的内存抽象。不过在Caffe2中,框架设计者可能是不想它太多余,于是将Serialization(从而将weights存成string)的功能给了它。
Anyway,这样我们在使用Caffe Operator时,会以Blob作为输入、输出(与Caffe一样),只是在Operator内部,一般需要使用Blob的data方法得到指向其元素的指针,然后再将它强制类型转换为合适的类型T(一般为Tensor),再使用它进行各种具体运算。
Caffe2中与Blob相关的代码如下。本节当中我们将重点介绍其中所涉及的Blob/BlobSerializaer/BlobStats等类及相关功能函数。
core]$ ls blob
blob_gpu_test.cc blob_serialization_gpu.cc blob_stats.cc
blob.h blob_serialization.h blob_stats.h
blob_serialization.cc blob_serializer_base.h blob_test.cc
Blob
以下为Blob的基本描述,可见看出它只有两个成员meta_与pointer_,分别表示指向存储对象的指针以及此指针的类型。
/**
* @brief Blob is a general container that hosts a typed pointer.
*
* A Blob hosts a pointer as well as its type, and takes charge of deleting it
* properly when the blob is deallocated or re-allocated with a new type. A blob
* could contain anything, although the most common case is to contain a Tensor.
*/
class CAFFE2_API Blob final {
public:
using DestroyCall = void(void*);
/**
* Initializes an empty Blob.
*/
Blob() : meta_(), pointer_(nullptr) {}
~Blob() { Reset(); }
Blob(Blob&& other) noexcept
: meta_(std::move(other.meta_)),
pointer_(std::move(other.pointer_)),
destroy_(std::move(other.destroy_)) {
other.meta_ = {};
other.pointer_ = nullptr;
other.destroy_ = nullptr;
}
........
........
};
通过下面两个成员函数,我们可以检查Blob所包含的对象的类型及是否是某种Device tensor类型等。
/**
* Checks if the content stored in the blob is of type T.
*/
template <class T>
bool IsType() const {
return meta_.Match<T>();
}
bool IsTensorType(DeviceType device_type) const {
bool is_match = meta_.Match<Tensor>();
auto* tensor = static_cast<Tensor*>(pointer_);
if (is_match && tensor && tensor->GetDeviceType() == device_type) {
return true;
}
return false;
}
以下为两种得到有类型指针与裸对象指针的办法。
/**
* @brief Gets the const reference of the stored object. The code checks if
* the stored object is of the desired type.
*/
// TODO(jerryzh): add a Get(DeviceType) function?
template <class T>
const T& Get() const {
CAFFE_ENFORCE(
IsType<T>(),
"wrong type for the Blob instance. Blob contains ",
meta_.name(),
" while caller expects ",
TypeMeta::TypeName<T>());
// TODO: after we add Get<Tensor>(DeviceType)
// and changed all the callsites, we can add
// a static assert here to enforce T != Tensor
return *static_cast<const T*>(pointer_);
}
const void* GetRaw() const {
return pointer_;
}
若想要对其存储对象进行写操作,则需要调用mutable_data方法,如下所示。
/**
* @brief Gets a mutable pointer to the stored object.
*
* If the current object is not of the right type, a new object is created
* and the old object is freed. Note that type T should have a default
* constructor. Otherwise, create the object yourself first, and use
* Reset().
*/
template <class T>
T* GetMutable() {
static_assert(
std::is_default_constructible<T>::value,
"GetMutable can't be called with non-default-constructible types. "
"Try using specialized methods");
static_assert(
!std::is_same<T, Tensor>::value,
"Use GetMutableTensor(DeviceType) instead");
if (IsType<T>()) {
return static_cast<T*>(pointer_);
} else {
VLOG(1) << "Create new mutable object " << TypeMeta::TypeName<T>();
return Reset<T>(new T());
}
}
inline Tensor* GetMutableTensor(DeviceType device_type) {
if (IsTensorType(device_type)) {
return static_cast<Tensor*>(pointer_);
} else {
VLOG(1) << "Create new mutable object " << TypeMeta::TypeName<Tensor>()
<< " DeviceType:" << device_type;
return Reset<Tensor>(new Tensor(device_type));
}
}
Reset 成员函数将使Blob得到此传入对象的ownership。在此前总要先释放指之前拥有的对象的ownership。
/**
* Sets the underlying object to the allocated one. The Blob then takes over
* the ownership of the passed in pointer. If there is already an object in
* the Blob, the old object is freed.
*
* This is used when the underlying class T does not have a default ctor, or
* complex initializations needs to be done outside the blob.
*/
template <class T>
T* Reset(T* allocated) {
if (pointer_ && destroy_) {
destroy_(pointer_);
}
meta_ = TypeMeta::Make<T>();
pointer_ = static_cast<void*>(allocated);
destroy_ = &Destroy<T>;
return allocated;
}
inline void*
Reset(void* allocated, const TypeMeta& meta, DestroyCall* destroy) {
if (pointer_ && destroy_) {
destroy_(pointer_);
}
meta_ = meta;
pointer_ = static_cast<void*>(allocated);
destroy_ = destroy;
return allocated;
}
/**
* Resets the Blob to an empty one.
*/
inline void Reset() {
if (pointer_ && destroy_) {
destroy_(pointer_);
}
pointer_ = nullptr;
meta_ = TypeMeta();
destroy_ = nullptr;
}
ShareExternal与Reset相反,它只享用此传入对象,但并不负责释放它即并不需对它付责任。
/**
* Sets the underlying object to the allocated one, but does not take over
* the ownership of the passed in pointer. If there is already an object in
* the Blob, the old object is freed.
*
* Unlike Reset, this does not take over the ownership of the pointer and the
* caller is responsible for making sure that the lifetime of the allocated
* blob outlasts the lifetime of any access to this blob, until another Reset
* call is made or the blob is destructed.
*/
template <class T>
typename std::remove_const<T>::type* ShareExternal(
typename std::remove_const<T>::type* allocated) {
return static_cast<T*>(ShareExternal(
static_cast<void*>(allocated),
TypeMeta::Make<typename std::remove_const<T>::type>()));
}
void* ShareExternal(void* allocated, const TypeMeta& meta) {
if (pointer_ && destroy_) {
destroy_(pointer_);
}
meta_ = meta;
pointer_ = static_cast<void*>(allocated);
destroy_ = nullptr;
return allocated;
}
Blob承担了部分Serialize的功能,可见所有的Weights需要放入到Blob里面正是需要仰仗它的这一功能来进行checkpoints存取。
/**
* Serializes the current 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.
*/
void Serialize(
const string& name,
BlobSerializerBase::SerializationAcceptor acceptor,
int chunk_size = kDefaultChunkSize) const;
/**
* @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.
*/
string Serialize(const string& name) const;
/**
* Deserializes from a string containing either BlobProto or TensorProto. If
* the deserialization fails, the content in the blob should no longer be
* trusted.
*/
void Deserialize(const string& content);
void Deserialize(const BlobProto& proto);
最后则为Blob私有空间的一些成员与公共函数。Destroy是一个static 模板成员函数,用在这里是再合适不过了。
private:
/**
* @brief A destroy call that is used to properly deconstruct objects.
*/
template <class T>
static void Destroy(void* pointer) {
delete static_cast<T*>(pointer);
}
TypeMeta meta_;
void* pointer_ = nullptr;
DestroyCall* destroy_ = nullptr;
AT_DISABLE_COPY_AND_ASSIGN(Blob);
};
BlobSerializerBase和BlobDeserializerBase
下面为BlobSerializerBase的概况,它是一个实现Blob serialization功能的虚基类。不同类型的Blob需要分别继承它来实现自己的Serialization操作。
/**
* @brief BlobSerializerBase is an abstract class that serializes a blob to a
* string.
*
* This class exists purely for the purpose of registering type-specific
* serialization code. If you need to serialize a specific type, you should
* write your own Serializer class, and then register it using
* REGISTER_BLOB_SERIALIZER. For a detailed example, see TensorSerializer for
* details.
*/
class BlobSerializerBase {
public:
virtual ~BlobSerializerBase() {}
using SerializationAcceptor =
std::function<void(const std::string& blobName, const std::string& data)>;
下面为它的两个主要的Serialization功能函数。
* @brief The virtual function that returns a serialized string for the input
* blob.
* @param blob
* the input blob to be serialized.
* @param name
* the blob name to be used in the serialization implementation. It is up
* to the implementation whether this name field is going to be used or
* not.
* @param acceptor
* a lambda which accepts key value pairs to save them to storage.
* serailizer can use it to save blob in several chunks
* acceptor should be thread-safe
*/
virtual void Serialize(const Blob& blob, const std::string& name,
SerializationAcceptor acceptor) = 0;
virtual void SerializeWithChunkSize(
const Blob& blob,
const std::string& name,
SerializationAcceptor acceptor,
int /*chunk_size*/) {
// Base implementation.
Serialize(blob, name, acceptor);
}
};
我们需要对每个类型Blob生成其特定的BlobSerializer子类。
// The Blob serialization registry and serializer creator functions.
CAFFE_DECLARE_TYPED_REGISTRY(
BlobSerializerRegistry,
TypeIdentifier,
BlobSerializerBase,
std::unique_ptr);
#define REGISTER_BLOB_SERIALIZER(id, ...) \
CAFFE_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);
}
相对应的有个Deserializer虚基类提供了Deserialization需要的函数接口。
/**
* @brief BlobDeserializerBase is an abstract class that deserializes a blob
* from a BlobProto or a TensorProto.
*/
class CAFFE2_API BlobDeserializerBase {
public:
virtual ~BlobDeserializerBase() {}
// Deserializes from a BlobProto object.
virtual void Deserialize(const BlobProto& proto, Blob* blob) = 0;
};
CAFFE_DECLARE_REGISTRY(BlobDeserializerRegistry, BlobDeserializerBase);
#define REGISTER_BLOB_DESERIALIZER(name, ...) \
CAFFE_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);
}
TensorSerializer和TensorDeserializer
TensorSerializer为BlobSerializerBase的一个子类,顾名思义,它主要用来实现Tensor类型的Serialization操作。同样还有一个为TensorDeserializer,它是BlobDeserializerBase的子类。
下面为在进行Serialization时的细节实现。可见主要是将需要的类型数据存到Protocol buffer里面,然后再使用它的功能来进行serialization/deserialization。
namespace detail {
template <typename SrcType, typename DstType>
inline void CopyToProtoAsIs(
const size_t size,
const SrcType* src,
google::protobuf::RepeatedField<DstType>* 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 (int i = 0; i < size; ++i) {
field->Add(0);
}
context->template CopyToCPU<SrcType>(
size, src, reinterpret_cast<SrcType*>(field->mutable_data()));
// Make sure that we finish the copy into the protobuf.
context->FinishDeviceComputation();
}
template <typename SrcType, typename DstType>
inline void CopyToProtoWithCast(
const size_t size,
const SrcType* src,
google::protobuf::RepeatedField<DstType>* 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<SrcType[]> buffer(new SrcType[size]);
context->template CopyToCPU<SrcType>(size, src, buffer.get());
context->FinishDeviceComputation();
field->Reserve(size);
for (int i = 0; i < size; ++i) {
field->Add(static_cast<DstType>(buffer[i]));
}
}
以下为Blob里面的两个Serialize函数实现,可以看出它主要是借助不同类型的BlobSerializer来完成此功能。Deserialize函数的实现与此类似,在此不再赘述。
// The blob serialization member function implementation.
void Blob::Serialize(
const string& name,
BlobSerializerBase::SerializationAcceptor acceptor,
int chunk_size) const {
std::unique_ptr<BlobSerializerBase> serializer(CreateSerializer(meta_.id()));
CAFFE_ENFORCE(serializer, "No known serializer for ", meta_.name());
serializer->SerializeWithChunkSize(*this, name, acceptor, chunk_size);
}
// The blob serialization member function implementation.
std::string Blob::Serialize(const string& name) const {
std::string data;
BlobSerializerBase::SerializationAcceptor acceptor = [&data](
const std::string&, const std::string& blob) {
DCHECK(data.empty()); // should be called once with kNoChunking
data = blob;
};
this->Serialize(name, acceptor, kNoChunking);
return data;
}
BlobStatGetter
Blob里面提供了些辅助类来提供些统计等功能如BlobStatGetter类。由下面代码可以看出,它亦是通过Type Id来选择使用不同的子类StatGetter的。
struct BlobStatGetter {
virtual size_t sizeBytes(const Blob& blob) const = 0;
virtual ~BlobStatGetter() {}
};
struct BlobStatRegistry {
private:
std::unordered_map<TypeIdentifier, std::unique_ptr<BlobStatGetter>> map_;
void doRegister(TypeIdentifier id, std::unique_ptr<BlobStatGetter>&& v);
public:
template <typename T, typename Getter>
struct Registrar {
Registrar() {
BlobStatRegistry::instance().doRegister(
TypeMeta::Id<T>(), std::unique_ptr<Getter>(new Getter));
}
};
const BlobStatGetter* get(TypeIdentifier id);
static BlobStatRegistry& instance();
};
以下为具体的对其使用。
const BlobStatGetter* BlobStatRegistry::get(TypeIdentifier id) {
auto it = map_.find(id);
if (it == map_.end()) {
return nullptr;
}
return it->second.get();
}
BlobStatRegistry& BlobStatRegistry::instance() {
static BlobStatRegistry registry;
return registry;
}
void BlobStatRegistry::doRegister(
TypeIdentifier id,
std::unique_ptr<BlobStatGetter>&& v) {
// don't use CAFFE_ENFORCE_EQ to avoid static initialization order fiasco.
if (map_.count(id) > 0) {
throw std::runtime_error("BlobStatRegistry: Type already registered.");
}
map_[id] = std::move(v);
}
namespace BlobStat {
size_t sizeBytes(const Blob& blob) {
auto* p = BlobStatRegistry::instance().get(blob.meta().id());
return p ? p->sizeBytes(blob) : 0;
}
} // namespace BlobStats
网友评论