介绍
Caffe2 core code中与tensor相关的可见于以下几个文件。
$ ls tensor tensor.cc tensor.h tensor_int8.cc tensor_int8.h
Tensor是Caffe2中的连续内存区域抽象表示。真正的caffe2 code中Tensor主要作为一个内存抽象APIs集合来供外部如Operator等对象来使用，其内部的大多数功能都实际依靠TensorImpl这个类来完成。TensorImpl的主要成员有两个，一个为dims_，它包含了当下内存的维度层次表示，另外一个则为storage_，它是一个Storage对象，亦是一个wrapper，实际起作用的为StorageImple，里面包含了此内存的实际地址，包含元素的类型(TypeMeta)等，同Tensor与TensorImpl的结构类似。
新一版Caffe2里面使用了许多Aten（有名的Tensor操作库，之前主要由Pytorch来使用）里的元素。像TensorImpl与Storage都是c10::intrusive_ptr_target的子类，内部自带了引用计数的操作。自然Tensor里在调用它们时都是通过使用类型为c10::intrusive_ptr的成员来做的。
下面我们将分别介绍下StorageImpl，Storage，TensorImpl及Tensor的内容。至于c10::intrusive_ptr_target与c10::intrusive_ptr多属于Aten的内容，在这里暂不作过多说明。
StorageImpl
新的Tensor实现里不再以Template的形式来支持不同类型的Context，而是通过将Tensor实现所需的DeviceType在构造时传入，并再转而构造一个Device specific的Storage对象来实现不同Device context支持。Storage里面的大多数功能都通过StorageImpl来实现。
下面的StorageImpl构造函数可以看出这一作法。
explicit StorageImpl(DeviceType device_type) : device_type_(device_type) {} StorageImpl(DeviceType device_type, TypeMeta data_type)
: data_type_(data_type), device_type_(device_type) {}当然更常用的是我们同时指出管理内存的实际地址、包含空间大小、上面元素的类型及对其进行删除的方法。如下所示：
template <typename Deleter = MemoryDeleter>
StorageImpl(
DeviceType device_type,
TypeMeta data_type, void* src, size_t capacity,
Deleter d = nullptr)
: data_type_(data_type), device_type_(device_type) {
CAFFE_ENFORCE_WITH_CALLER(
data_type_.id() != TypeIdentifier::uninitialized(), "To create storage with a raw external pointer you need to pass in an "
"initialized data_type(TypeMeta)."); // Check if the deleter is a MemoryDeleter and is a simple nullptr.
if (std::is_same<MemoryDeleter, Deleter>::value && reinterpret_cast<MemoryDeleter*>(static_cast<void*>(&d))[0] == nullptr) { // Use aliasing constructor trick to avoid calling the destructor.
data_ptr_ = std::shared_ptr<void>(std::shared_ptr<void>(), src);
} else {
data_ptr_.reset(src, d);
}
capacity_ = capacity;
}明白了StorageImpl里面所含有的基本成员，我们就清楚了它所能提供的对这些成员进行操作的一些方法，举例如下：
...........
........... void reset() {
data_ptr_.reset();
capacity_ = 0;
} template <typename T> inline bool IsType() const { return data_type_.Match<T>();
} void* data() const { return data_ptr_.get();
} void* data() { return data_ptr_.get();
} DataPtr& data_ptr() { return data_ptr_;
}
...........
...........Storage
Storage对象里面包含一个类型为c10::instrusitive_ptr的StorageImpl成员。然后它以外包的形式，向外开放StorageImpl所提供的一些功能支持。
如下为Storage的一个较全的构造函数。
template <typename Deleter = MemoryDeleter>
Storage( void* src,
DeviceType device_type,
TypeMeta data_type, size_t capacity,
Deleter d = nullptr)
: storage_impl_(c10::make_intrusive<StorageImpl>(
device_type,
data_type,
src,
capacity,
d)) {}然后通过外包，向外提供一些基本的内存单元属性读取或设置功能。
...............
............... void reset() {
storage_impl_->reset();
} template <typename T> inline bool IsType() const { return storage_impl_->IsType<T>();
} void* data() const { return storage_impl_->data();
} void* data() { return storage_impl_->data();
} DataPtr& data_ptr() { return storage_impl_->data_ptr();
} const DataPtr& data_ptr() const { return storage_impl_->data_ptr();
}
...............
............... inline long use_count() const { return storage_impl_.use_count();
} inline bool unique() const { return storage_impl_.unique();
} template <typename Deleter = MemoryDeleter> void UniqueStorageShareExternalPointer( void* src, const DataType& data_type, size_t capacity,
Deleter d = nullptr) {
CAFFE_ENFORCE_WITH_CALLER(
storage_impl_.unique(), "UniqueStorageShareExternalPointer can only be called when \
use_count == 1");
storage_impl_->UniqueStorageShareExternalPointer<Deleter>(
src, data_type, capacity, d);
} protected:
c10::intrusive_ptr<StorageImpl> storage_impl_;
};TensorImpl
TensorImpl对内存的管理通过两个成员完成，一个为dims，另一个则为storage_。其中storage_为上面讲过的一个Storage对象，里面有所管理内存的实际地址，空间大小，类型(TypeMeta)等。
class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target {
public:
TensorImpl() = delete; explicit TensorImpl(DeviceType device_type) : storage_(device_type) {} /**
* @brief Creates a tensor of the given dimension.
*
* Note that the actual data allocation is not going to be carried out until
* the first time mutable_data() is called.
*/
// TODO: here, we create a Storage
// and immediately discard it in Resize() since
// reset_tensor will be true and FreeMemory will be called,
// we might want to avoid creating Storage twice?
explicit TensorImpl(const vector<TIndex>& dims, at::DeviceType device_type)
: storage_(device_type) {
Resize(dims);
}我们可以move copy或assign tensorImpl对象，但却不可以以复制copy的形式进行操作。
/**
* @brief Delete the copy constructor and use Clone explicitly
*/
TensorImpl(const TensorImpl& src) = delete;
TensorImpl(TensorImpl&& src) noexcept {
swap(src);
}
TensorImpl& operator=(TensorImpl&&) = default; // Note(jiayq): possibly a rule-of-three violation, but we explicitly
// discourage the use of = for Tensors.
TensorImpl& operator=(const TensorImpl& src) = delete;因为Tensor去掉了Context的模板参数，因此将它作为一个static的成员放在了类里面，指向DeviceType所对应的Context（DeviceType则存在Storage对象成员里面如上节所讲。）
/*
* Since we removed template from tensor, we now store a static
* context pointer in tensor, which indicates the type of the tensor.
*/
BaseStaticContext* GetStaticContext() const { return get_static_context(GetDeviceType());
} /* @brief
* Create a context that has the same device_type
* as the tensor.
* Note that this doesn't support passing in argument
* TODO(jerryzh): move this to a global registry
* that can create context for us
*/
std::unique_ptr<BaseContext> CreateContext() const { return GetStaticContext()->CreateContext();
}如下为一个Tensor扩展其内存空间的方法。
/**
* @brief Extends the outer-most dimension of this tensor by num elements,
* preserving the existing data.
*
* The underlying data may be reallocated in order to accommodate the new
* elements, in which case this tensors' capacity is grown at a factor of
* growthPct. This ensures that Extend runs on an amortized O(1) time
* complexity.
*/
void Extend(TIndex num, float growthPct, BaseContext* context) {
CAFFE_ENFORCE_WITH_CALLER(
is_contiguous_, "Tensor must be contiguous in order to call Extend.");
CAFFE_ENFORCE_GE_WITH_CALLER(dims_.size(), 1);
CAFFE_ENFORCE_GE_WITH_CALLER(
num, 0, "`num` must be non-negative for Extend"); auto newDims = dims_;
newDims[0] += num; if (!storage_.data()) {
Resize(newDims); return;
} auto newNumel = std::accumulate(
newDims.begin(),
newDims.end(), static_cast<TIndex>(1), std::multiplies<TIndex>()); if (newNumel * storage_.itemsize() <= storage_.capacity()) {
dims_ = newDims;
numel_ = newNumel; return;
} auto newCapacity = dims_;
newCapacity[0] = std::max<size_t>(
newDims[0], std::ceil(dims_[0] * (growthPct + 100) / 100)); auto oldData = std::move(storage_.data_ptr()); auto oldSize = numel_; auto oldDims = dims_;
Resize(newCapacity); auto* newData = raw_mutable_data(storage_.dtype());
CAFFE_ENFORCE(
context != nullptr, "Context must be provided to Extend the tensor");
context->CopyItemsSameDevice(
storage_.dtype(), oldSize, oldData.get(), newData);
reserved_ = true;
dims_ = newDims;
numel_ = newNumel;
}以下为对Tensor管理空间进行shrink时所做的事。注意我们不可对共享的storage单元进行shrink操作。
/**
* @brief Shrinks the outer-most dimension to given size, keeping the data.
*
* This method guarantees that no re-allocations are carried out, which means
* that the extra capacity after the end of the shurnk tensor is maintained.
*/
void ShrinkTo(TIndex outer_dim) {
CAFFE_ENFORCE_WITH_CALLER(
is_contiguous_, "Tensor must be contiguous in order to call ShrinkTo.");
CAFFE_ENFORCE_WITH_CALLER(dims_.size() >= 1, "Tensor must be at least 1D");
CAFFE_ENFORCE_WITH_CALLER(
outer_dim <= dims_[0], "New outer dimension must be smaller than current.");
CAFFE_ENFORCE(
storage_.unique(), "Can't call ShrinkTo on shared storage, please call Resize instead.");
dims_[0] = outer_dim;
numel_ = std::accumulate(
dims_.begin(),
dims_.end(), static_cast<TIndex>(1), std::multiplies<TIndex>());
}Tensor上的Resize操作，则主要是更改dims；只有某些情况下，一些flags存在，加上实时条件满足则会真正地释放掉旧的内存，下一次mutable_data调用时则重新分配内存。
/**
* @brief Resizes a tensor.
*
* Resize takes in a vector of ints specifying the dimensions of the tensor.
* You can pass in an empty vector to specify that it is a scalar (i.e.
* containing one single item).
*
* The underlying storage may be deleted after calling Resize: if the new
* shape leads to a different number of items in the tensor, the old memory
* is deleted and new memory will be allocated next time you call
* mutable_data(). However, if the shape is different but the total number of
* items is the same, the underlying storage is kept.
*/
template <typename... Ts> void Resize(Ts... dim_source) { bool is_init = numel_ == -1; bool size_changed = SetDims(dim_source...); if (size_changed) { // If needed, we will free the data. the next mutable_data() call
// will create the data storage.
bool reset_tensor = false; if (reserved_) { // If tensor is reserved then don't claim its memeory unless capacity()
// is smaller than new size
reset_tensor = storage_.capacity() < numel_ * storage_.itemsize();
} else {
reset_tensor = storage_.capacity() < numel_ * storage_.itemsize() ||
!FLAGS_caffe2_keep_on_shrink ||
storage_.capacity() - numel_ * storage_.itemsize() >
FLAGS_caffe2_max_keep_on_shrink_memory;
} if (reset_tensor && !is_init) {
FreeMemory();
}
}
}Reshape则只是需要保证新的dims与旧的dims总的size相同，然后直接改变dims，而对底下的storage对象则并不改动。
/**
* Resizes the tensor without touching underlying storage.
* This requires the total size of the tensor to remains constant.
*/
inline void Reshape(const vector<TIndex>& dims) {
CAFFE_ENFORCE_WITH_CALLER(
is_contiguous_, "Tensor must be contiguous in order to call Reshape.");
TIndex new_size = 1; for (auto d : dims) {
CAFFE_ENFORCE_GE_WITH_CALLER(d, 0);
new_size *= d;
}
CAFFE_ENFORCE_WITH_CALLER(
new_size == numel_, "New size and old size are not equal. You cannot use Reshape, "
"but should use Resize."
// TODO(jiayq): remove the following warning after pending diffs
// stabilize.
" The old caffe2 mixes Reshape and Resize but this behavior has "
"been changed. If you find this error, most likely you will need "
"to change corresponding code from Reshape to Resize.");
dims_ = dims;
}ShareData则是为了在多个Tensor之间共享底部的Storage，而它们可以有着完全不同的dims，只需要保证其size相同即可。
/**
* @brief Shares the data with another tensor.
*
* To share data between two tensors, the sizes of the two tensors must be
* equal already. The reason we do not implicitly do a Resize to make the two
* tensors have the same shape is that we want to allow tensors of different
* shapes but the same number of items to still be able to share data. This
* allows one to e.g. have a n-dimensional Tensor and a flattened version
* sharing the same underlying storage.
*
* The source tensor should already have its data allocated.
*/
void ShareData(const TensorImpl& src) { // Right now, we are assuming the device_type are the same, since it is
// inherently the same in the non-templatized code. We should probably add
// an ENFORCE here which might affect perf a little bit.
CAFFE_ENFORCE_EQ_WITH_CALLER(
src.numel_,
numel_, "Size mismatch - did you call reshape before sharing the data?"); // It is possible that the source tensor hasn't called mutable_data() yet,
// in which case ShareData() doesn't make much sense since we don't really
// know what to share yet.
CAFFE_ENFORCE_WITH_CALLER(
src.storage_.data() || src.numel_ == 0, "Source tensor has no content and has size > 0"); // Finally, do sharing.
/* Since we create new Storage whenever we need to change data_type/capacity
* this still keeps the original semantics
*/
storage_ = src.storage();
}以下为返回只读的raw data或者带类型（Typename T）的data的方法。
/**
* Returns a const raw void* pointer of the underlying storage. mutable_data()
* or raw_mutable_data() must have been called prior to this function call.
*/
inline const void* raw_data() const {
CAFFE_ENFORCE_WITH_CALLER(
is_contiguous_, "Tensor must be contiguous in order to call raw_data()");
CAFFE_ENFORCE_WITH_CALLER(storage_.data() || numel_ == 0); return storage_.data();
} /**
* Returns a typed pointer of the underlying storage. mutable_data() or
* raw_mutable_data() must have been called prior to this function call, and
* the data type must be of the correct type. If you want to get a void*
* pointer instead, use raw_data().
*/
template <typename T> inline const T* data() const {
CAFFE_ENFORCE_WITH_CALLER(
is_contiguous_, "Tensor must be contiguous in order to call data()");
CAFFE_ENFORCE_WITH_CALLER(
storage_.data() || numel_ == 0, "The tensor is of non-zero shape, but its data is not allocated yet. "
"Caffe2 uses a lazy allocation, so you will need to call "
"mutable_data() or raw_mutable_data() to actually allocate memory.");
CAFFE_ENFORCE_WITH_CALLER(
IsType<T>(), "Tensor type mismatch, caller expects elements to be ",
TypeMeta::TypeName<T>(), " while tensor contains ",
storage_.dtype().name()); return static_cast<T*>(storage_.data());
}以下为有意思的mutable_data的实现方法。其实里面有一个delay执行分配内存的优化做法。
/**
* Returns a mutable raw pointer of the underlying storage. Since we will need
* to know the type of the data for allocation, a TypeMeta object is passed in
* to specify the necessary information. This is conceptually equivalent of
* calling mutable_data<T>() where the TypeMeta parameter meta is derived from
* the type T. This function differs from mutable_data<T>() in the sense that
* the type T can be specified during runtime via the TypeMeta object.
*
* If the existing data does not match the desired type, it will be deleted
* and a new storage will be created.
*/
inline void* raw_mutable_data(const TypeMeta& meta) {
CAFFE_ENFORCE_WITH_CALLER(
is_contiguous_, "Tensor must be contiguous in order to call raw_mutable_data()"); // For 0-size tensors it's fine to return any pointer (including nullptr)
if (storage_.dtype() == meta && (storage_.data() || numel_ == 0)) { return storage_.data();
} else { bool had_special_dtor = storage_.dtype().dtor() != nullptr; if (storage_.unique()) {
storage_.set_dtype(meta); // TODO: recalcuate numel when we store numel instead of capacity in
// Storage
} else { if (storage_.dtype() != meta) {
storage_ = Storage(storage_.device_type(), meta);
}
}
CAFFE_ENFORCE_WITH_CALLER(
numel_ >= 0, "Tensor is not initialized. You probably need to call Resize() "
"before calling mutable_data()"); // We can reuse the existing buffer if the current data does not have
// a special destructor and the new data doesn't have a special
// constructor.
if (numel_ == 0 ||
(meta.ctor() == nullptr && !had_special_dtor &&
storage_.capacity() >= numel_ * storage_.itemsize())) { return storage_.data();
} if (meta.ctor()) { // For types that need placement new, we will call it, as well as
// making sure that when the data is freed, it calls the right
// destruction procedure.
auto size = numel_; auto dtor = storage_.dtype().dtor(); auto ptr_and_deleter =
GetStaticContext()->New(numel_ * storage_.itemsize()); auto deleter = ptr_and_deleter.second;
storage_.data_ptr().reset(
ptr_and_deleter.first, [size, dtor, deleter](void* ptr) -> void {
dtor(ptr, size);
deleter(ptr);
});
storage_.dtype().ctor()(storage_.data(), numel_);
} else { // For fundamental type, new and delete is easier.
auto ptr_and_deleter =
GetStaticContext()->New(numel_ * storage_.itemsize());
storage_.data_ptr().reset(
ptr_and_deleter.first, ptr_and_deleter.second);
}
storage_.set_numel(numel_); return storage_.data();
}
}下面为它的一些protected成员。
protected: DimVector dims_; // sizes_ DimVector strides_; TIndex numel_ = -1; // numel_ bool is_contiguous_ = true; // we decide to keep reserved_ and it will // live in Tensor after the split // The logic is that if Extend() or ReserveSpace() were ever called, // then subsequent Resize()s will not free up Storage. bool reserved_ = false; Storage storage_; // int64_t storage_offset_;
我们还有一个static的UndefinedTensor使用。
class CAFFE2_API UndefinedTensorImpl final : public TensorImpl {
UndefinedTensorImpl() : TensorImpl(CPU){}; public: // Without this, we get:
// error: identifier "at::UndefinedTensor::_singleton" is undefined in device code
// (ostensibly because the constexpr tricks MSVC into trying to compile this
// function for device as well).#ifdef _WIN32
static inline TensorImpl * singleton() {#else
static constexpr inline TensorImpl * singleton() {#endif
return &singleton_;
} private: static UndefinedTensorImpl singleton_;
};Tensor
最后我们来看下真正外部程序所见的对象，Tensor。
以下为它的构造函数及基本类成员，从此易知它的大部分操作都是借助TensorImpl来完成的。在它的多个构造函数当中还有一个模板构造函数。
/**
* @brief Tensor class holds a shared pointer to the implementation TensorImpl,
* redirects API calls to TensorImpl;
* Copying of Tensor results in sharing the same underlying implementation
* object
*/class CAFFE2_API Tensor final { protected:
using TensorImplPtr = c10::intrusive_ptr<TensorImpl, UndefinedTensorImpl>;
TensorImplPtr impl_; public:
Tensor() : impl_() {}
operator bool() const { return impl_.defined();
}
explicit Tensor(const vector<TIndex>& dims, DeviceType type)
: impl_(
c10::make_intrusive<TensorImpl, UndefinedTensorImpl>(dims, type)) {}
template <
typename T,
typename = typename std::enable_if<std::is_scalar<T>::value>::type>
Tensor(const T& value, BaseContext* context)
: impl_(c10::make_intrusive<TensorImpl, UndefinedTensorImpl>(
value,
context)) {}以下为Tensor借助TensorImpl实现一些操作的提现。
inline int ndim() const { return impl_.get()->ndim();
} inline TIndex size() const { return impl_.get()->size();
} inline size_t itemsize() const { return impl_.get()->itemsize();
} inline size_t nbytes() const { return impl_.get()->nbytes();
} inline size_t capacity_nbytes() const { return impl_.get()->capacity_nbytes();
}
作者：manofmountain
链接：https://www.jianshu.com/p/21c74798062b
随时随地看视频
