| /*! |
| * Copyright (c) 2014 by Contributors |
| * \file tensor_blob.h |
| * \brief TBlob class that holds common representation of |
| * arbirary dimension tensor, can be used to transformed |
| * to normal fixed dimenson tensor |
| * \author Tianqi Chen |
| */ |
| #ifndef MXNET_TENSOR_BLOB_H_ |
| #define MXNET_TENSOR_BLOB_H_ |
| |
| #include <dmlc/logging.h> |
| #include <dmlc/json.h> |
| #include <dlpack/dlpack.h> |
| #include <vector> |
| #include <iostream> |
| #include <utility> |
| #include <algorithm> |
| #include "./base.h" |
| #if MXNET_USE_MKL2017 == 1 |
| #include <mkl_memory.h> |
| #endif |
| namespace mxnet { |
| |
| /* Forward declaration for friend declaration in TBlob */ |
| class NDArray; |
| |
| /*! |
| * \brief tensor blob class that can be used to hold tensor of any dimension, |
| * any device and any data type, |
| * This is a weak type that can be used to transfer data through interface |
| * TBlob itself do not involve any arithmentic operations, |
| * but it can be converted to tensor of fixed dimension for further operations |
| * |
| * Like tensor, this data structure is like a pointer class and do not |
| * implicit allocated, de-allocate space. |
| * This data structure can be helpful to hold tensors of different dimensions |
| * and wait for further processing |
| */ |
| class TBlob { |
| friend class NDArray; |
| public: |
| /*! \brief pointer to the data */ |
| void *dptr_; |
| /*! \brief shape of the tensor */ |
| TShape shape_; |
| /*! \brief type flag of the tensor blob */ |
| int type_flag_; |
| |
| /*! \brief storing mkl chunk buffer blob, use for experimental only */ |
| #if MKL_EXPERIMENTAL == 1 |
| std::shared_ptr<MKLMemHolder> Mkl_mem_; |
| #endif |
| /*! \brief default constructor, default copy assign will work */ |
| TBlob(void) |
| : dptr_(NULL), |
| type_flag_(mshadow::DataType<real_t>::kFlag) { |
| #if MKL_EXPERIMENTAL == 1 |
| Mkl_mem_ = NULL; |
| #endif |
| SetDLTensor(cpu::kDevMask, 0); |
| } |
| /*! |
| * \brief constructor that construct TBlob from contiguous memory |
| * \param dptr the pointer to the memory |
| * \param shape the shape of the data |
| * \param dev_mask the device mask, can be cpu::kDevMask or gpu::kDevMask |
| * \param dev_id the device id |
| */ |
| template<typename DType> |
| TBlob(DType *dptr, const TShape &shape, int dev_mask, int dev_id = -1) |
| : dptr_(dptr), shape_(shape), |
| type_flag_(mshadow::DataType<DType>::kFlag) { |
| #if MKL_EXPERIMENTAL == 1 |
| Mkl_mem_ = NULL; |
| #endif |
| SetDLTensor(dev_mask, dev_id); |
| } |
| /*! |
| * \brief constructor that construct TBlob from contiguous memory |
| * \param dptr the pointer to the memory |
| * \param shape the shape of the data |
| * \param dev_mask the device mask, can be cpu::kDevMask or gpu::kDevMask |
| * \param type_flag the type flag. Can be one of enum mshadow::dtype |
| * \param dev_id the device id |
| */ |
| TBlob(void *dptr, const TShape &shape, int dev_mask, int type_flag, int dev_id = -1) |
| : dptr_(dptr), shape_(shape), type_flag_(type_flag) { |
| #if MKL_EXPERIMENTAL == 1 |
| Mkl_mem_ = NULL; |
| #endif |
| SetDLTensor(dev_mask, dev_id); |
| } |
| /*! |
| * \brief constructor from tensor |
| * \param src source tensor |
| * \tparam Device which device the tensor is on |
| * \tparam dim tensor dimension |
| * \tparam DType the type of elements in the tensor |
| */ |
| template<typename Device, int dim, typename DType> |
| TBlob(const mshadow::Tensor<Device, dim, DType> &src) { // NOLINT(*) |
| *this = src; |
| } |
| /*! |
| * \brief assignment from tensor |
| * \param src source tensor |
| * \tparam Device which device the tensor is on |
| * \tparam dim tensor dimension |
| * \tparam DType the type of elements in the tensor |
| * \return reference of self |
| */ |
| template<typename Device, int dim, typename DType> |
| inline TBlob &operator=(const mshadow::Tensor<Device, dim, DType> &src) { |
| dptr_ = src.dptr_; |
| shape_ = src.shape_; |
| type_flag_ = mshadow::DataType<DType>::kFlag; |
| SetDLTensor(Device::kDevMask, -1); |
| #if MKL_EXPERIMENTAL == 1 |
| Mkl_mem_ = NULL; |
| #endif |
| return *this; |
| } |
| /*! |
| * \return whether the tensor's memory is continuous |
| */ |
| inline bool CheckContiguous(void) const { |
| return true; |
| } |
| /*! |
| * \brief reshape to shape |
| * \param shape desired shape |
| * \return reshaped blob |
| */ |
| inline TBlob reshape(const TShape& shape) const { |
| CHECK_EQ(this->shape_.Size(), shape.Size()) << "Shape size mismatch " |
| << this->shape_.Size() << " v.s. " << shape.Size(); |
| TBlob ret(this->dptr_, shape, this->dev_mask(), this->type_flag_, this->dev_id()); |
| return ret; |
| } |
| /*! |
| * \brief flatten the tensor to 2 dimension, collapse the higher dimensions together |
| * \param stream the possible stream target tensor should reside on |
| * \tparam Device which device the tensor is on |
| * \tparam DType the type of elements in the tensor |
| * \return tensor after flatten |
| */ |
| template<typename Device, typename DType> |
| inline mshadow::Tensor<Device, 2, DType> FlatTo2D( |
| mshadow::Stream<Device> *stream = NULL) const { |
| CHECK(Device::kDevMask == this->dev_mask()) |
| << "TBlob.get: device type do not match specified type"; |
| CHECK(mshadow::DataType<DType>::kFlag == type_flag_) |
| << "TBlob.get_with_shape: data type do not match specified type." |
| << "Expected: " << type_flag_ << " v.s. given " << mshadow::DataType<DType>::kFlag; |
| #if MKL_EXPERIMENTAL == 1 |
| if (Mkl_mem_ != nullptr) { |
| Mkl_mem_->check_and_prv_to_cpu(dptr_); |
| } |
| #endif |
| return mshadow::Tensor<Device, 2, DType>(static_cast<DType*>(dptr_), |
| shape_.FlatTo2D(), |
| shape_[shape_.ndim() - 1], |
| stream); |
| } |
| /*! |
| * \brief flatten the tensor to 1 dimension, collapse all the dimensions together. |
| * \param stream the possible stream target tensor should reside on |
| * \tparam Device which device the tensor is on |
| * \tparam DType the type of elements in the tensor |
| * \return tensor after flatten |
| */ |
| template<typename Device, typename DType> |
| inline mshadow::Tensor<Device, 1, DType> FlatTo1D( |
| mshadow::Stream<Device> *stream = NULL) const { |
| return this->get_with_shape<Device, 1, DType>( |
| mshadow::Shape1(shape_.Size()), stream); |
| } |
| /*! \brief return number of dimension of the tensor inside */ |
| inline int ndim(void) const { |
| return shape_.ndim(); |
| } |
| /*! |
| * \brief return size of i-th dimension, start counting from highest dimension |
| * \param idx the dimension count from the highest dimensin |
| * \return the size |
| */ |
| inline index_t size(index_t idx) const { |
| return shape_[idx]; |
| } |
| /*! \brief total number of elements in the tensor */ |
| inline index_t Size(void) const { |
| return shape_.Size(); |
| } |
| /*! \brief get pointer in dtype */ |
| template<typename DType> |
| inline DType* dptr() const { |
| CHECK(mshadow::DataType<DType>::kFlag == type_flag_) |
| << "TBlob.get_with_shape: data type do not match specified type." |
| << "Expected: " << type_flag_ << " v.s. given " << mshadow::DataType<DType>::kFlag; |
| #if MKL_EXPERIMENTAL == 1 |
| if (Mkl_mem_ != nullptr) { |
| Mkl_mem_->check_and_prv_to_cpu(dptr_); |
| } |
| #endif |
| return static_cast<DType*>(dptr_); |
| } |
| /*! \brief device mask of the corresponding device */ |
| inline int dev_mask() const { |
| return dltensor_.ctx.device_type; |
| } |
| /*! \brief device index of the corresponding device */ |
| inline int dev_id() const { |
| return dltensor_.ctx.device_id; |
| } |
| /*! |
| * \brief return the corresponding DLTensor |
| * \return the address of internal DLTensor |
| */ |
| inline const DLTensor& dltensor() const { |
| return dltensor_; |
| } |
| |
| /*! |
| * \brief fetch the tensor, with respect to specific dimension |
| * if dim do not match the stored dimension, an error will be issued |
| * \return the tensor requested |
| * \param stream the possible stream target tensor should reside on |
| * \tparam Device which device the tensor is on |
| * \tparam dim dimension of the tensor |
| * \tparam DType the type of elements in the tensor |
| */ |
| template<typename Device, int dim, typename DType> |
| inline mshadow::Tensor<Device, dim, DType> get(mshadow::Stream<Device> *stream = NULL) const { |
| CHECK(Device::kDevMask == this->dev_mask()) |
| << "TBlob.get: device type do not match specified type"; |
| return mshadow::Tensor<Device, dim, DType>(dptr<DType>(), |
| shape_.get<dim>(), shape_[shape_.ndim() - 1], stream); |
| } |
| /*! |
| * \brief fetch a tensor in given shape |
| * If size do not match the stored size, an error will be issued |
| * \return the tensor requested |
| * \param shape the shape required |
| * \param stream the possible stream target tensor should reside on |
| * \tparam Device which device the tensor is on |
| * \tparam dim dimension of the tensor |
| * \tparam DType the type of elements in the tensor |
| */ |
| template<typename Device, int dim, typename DType> |
| inline mshadow::Tensor<Device, dim, DType> get_with_shape( |
| const mshadow::Shape<dim> &shape, |
| mshadow::Stream<Device> *stream = NULL) const { |
| CHECK(Device::kDevMask == this->dev_mask()) |
| << "TBlob.get: device type do not match specified type"; |
| CHECK_EQ(this->CheckContiguous(), true) << "TBlob.get_reshape: must be contiguous"; |
| CHECK_EQ(this->shape_.Size(), shape.Size()) |
| << "TBlob.get_with_shape: new and old shape do not match total elements"; |
| return mshadow::Tensor<Device, dim, DType>(dptr<DType>(), shape, |
| shape[dim - 1], stream); |
| } |
| /*! |
| * \brief flatten the tensor to 3 dimension, |
| * collapse the dimension before and after specified axis. |
| * \param axis The axis specified. |
| * \param stream the possible stream target tensor should reside on |
| * \tparam Device which device the tensor is on |
| * \tparam DType the type of elements in the tensor |
| * \return tensor after flatten |
| */ |
| template<typename Device, typename DType> |
| inline mshadow::Tensor<Device, 3, DType> FlatTo3D( |
| int axis, mshadow::Stream<Device> *stream = NULL) const { |
| return this->get_with_shape<Device, 3, DType>( |
| this->shape_.FlatTo3D(axis), stream); |
| } |
| /*! |
| * \brief flatten the tensor to 3 dimension, |
| * collapse the dimension: [0, axis_begin), [axis_begin, axis_end], (axis_end, ndim). |
| * \param axis_begin The beginning axis specified. |
| * \param axis_end The ending axis specified. |
| * \param stream the possible stream target tensor should reside on |
| * \tparam Device which device the tensor is on |
| * \tparam DType the type of elements in the tensor |
| * \return tensor after flatten |
| */ |
| template<typename Device, typename DType> |
| inline mshadow::Tensor<Device, 3, DType> FlatTo3D( |
| int axis_begin, int axis_end, |
| mshadow::Stream<Device> *stream = NULL) const { |
| return this->get_with_shape<Device, 3, DType>( |
| this->shape_.FlatTo3D(axis_begin, axis_end), stream); |
| } |
| /*! |
| * \brief flatten the tensor to specified number of dimensions, |
| * collapse the highest dimensions or pad with higher dimensions |
| * \param stream the possible stream target tensor should reside on |
| * \tparam Device which device the tensor is on |
| * \tparam dim desired number of dimensions of returned tensor |
| * \tparam DType the type of elements in the tensor |
| * \return tensor after flatten |
| */ |
| template<typename Device, int dim, typename DType> |
| inline mshadow::Tensor<Device, dim, DType> FlatToKD( |
| mshadow::Stream<Device> *stream = NULL) const { |
| mshadow::Shape<dim> shape; |
| shape[0] = 1; |
| // Pad higher dimensions in case dim > ndim() |
| for (int i = 0; i < dim - ndim(); ++i) { |
| shape[i] = 1; |
| } |
| // Collapse higher dimensions in case dim < ndim() |
| for (int i = 0; i < ndim() - dim + 1; ++i) { |
| shape[0] *= shape_[i]; |
| } |
| // Preserve lower dimensions. |
| for (int i = std::max(0, ndim() - dim + 1); i < ndim(); ++i) { |
| shape[i - ndim() + dim] = shape_[i]; |
| } |
| return this->get_with_shape<Device, dim, DType>(shape, stream); |
| } |
| |
| private: |
| static DLDataType DTypeTransform(int type_flag) { |
| static std::unordered_map<int, DLDataType> |
| MSHADOW_DTYPE_TO_DLPACK_DTYPE = { |
| {0, {2, 32, 1}}, // Float32 |
| {1, {2, 64, 1}}, // Float64 |
| {2, {2, 16, 1}}, // Float16 |
| {3, {1, 8, 1}}, // UInt8 |
| {4, {0, 32, 1}}, // Int32 |
| {5, {0, 8, 1}} // Int8 |
| }; |
| return MSHADOW_DTYPE_TO_DLPACK_DTYPE[type_flag]; |
| } |
| |
| inline void SetDLTensor(int dev_mask, int dev_id) { |
| dltensor_.data = dptr_; |
| dltensor_.ctx = DLContext{static_cast<DLDeviceType>(dev_mask), dev_id}; |
| dltensor_.ndim = shape_.ndim(); |
| dltensor_.dtype = DTypeTransform(type_flag_); |
| dltensor_.shape = shape_.data(); |
| dltensor_.strides = NULL; |
| dltensor_.byte_offset = 0; |
| } |
| |
| private: |
| /*! \brief corresponding DLTensor of this TBlob */ |
| DLTensor dltensor_; |
| }; |
| } // namespace mxnet |
| |
| namespace dmlc { |
| // Add a few patches to support TShape in dmlc/parameter. |
| DMLC_DECLARE_TYPE_NAME(mxnet::TShape, "Shape(tuple)"); |
| DMLC_DECLARE_TYPE_NAME(nnvm::Tuple<int>, "Shape(tuple)"); |
| DMLC_DECLARE_TYPE_NAME(nnvm::Tuple<dmlc::optional<int>>, "Shape(tuple)"); |
| |
| namespace parameter { |
| |
| template<> |
| class FieldEntry<mxnet::TShape> |
| : public FieldEntryBase<FieldEntry<mxnet::TShape>, mxnet::TShape> { |
| public: |
| FieldEntry() : enforce_nonzero_(false), expect_ndim_(0) {} |
| // parent class |
| typedef FieldEntryBase<FieldEntry<mxnet::TShape>, mxnet::TShape> Parent; |
| |
| virtual void Check(void *head) const { |
| Parent::Check(head); |
| mxnet::TShape &v = this->Get(head); |
| if (expect_ndim_ != 0 && v.ndim() != expect_ndim_) { |
| std::ostringstream os; |
| os << "value " << v << "for Parameter " << this->key_ |
| << " has wrong dimensions, expected dimension=" << expect_ndim_; |
| throw dmlc::ParamError(os.str()); |
| } |
| if (enforce_nonzero_) { |
| for (mxnet::index_t i = 0; i < v.ndim(); ++i) { |
| if (v[i] == 0U) { |
| std::ostringstream os; |
| os << "value " << v << "for Parameter " << this->key_ |
| << " is invalid, the input shape must be nonzero in all dimensions"; |
| throw dmlc::ParamError(os.str()); |
| } |
| } |
| } |
| } |
| inline FieldEntry<mxnet::TShape> &enforce_nonzero() { |
| this->enforce_nonzero_ = true; |
| return this->self(); |
| } |
| inline FieldEntry<mxnet::TShape> &set_expect_ndim(mxnet::index_t ndim) { |
| expect_ndim_ = ndim; |
| return this->self(); |
| } |
| |
| private: |
| // whether all the entries need to be nonzero |
| bool enforce_nonzero_; |
| // expected number of dimension, default = 0 means no restriction. |
| mxnet::index_t expect_ndim_; |
| }; |
| |
| } // namespace parameter |
| } // namespace dmlc |
| |
| #endif // MXNET_TENSOR_BLOB_H_ |
