tensor_container.h

Go to the documentation of this file.

#ifndef MSHADOW_TENSOR_CONTAINER_H_
#define MSHADOW_TENSOR_CONTAINER_H_
#include "./tensor.h"
#include "./io.h"

namespace mshadow {
template<typename Device, int dimension, typename DType = default_real_t>
class TensorContainer: public Tensor<Device, dimension, DType> {
 public:
  explicit TensorContainer(bool pad = MSHADOW_ALLOC_PAD) {
    this->pad_ = pad;
    this->dptr_ = data_.dptr_ = NULL;
    this->shape_[0] = 0;
    this->stride_ = 0;
    this->data_.stride_ = 0;
    this->data_.shape_[0] = 0;
  }
  explicit TensorContainer(const Shape<dimension> &shape) {
    this->pad_ = MSHADOW_ALLOC_PAD;
    data_.dptr_ = NULL;
    this->AllocByShape(shape);
  }
  explicit TensorContainer(const Shape<dimension> &shape, DType initv) {
    this->pad_ = MSHADOW_ALLOC_PAD;
    data_.dptr_ = NULL;
    this->AllocByShape(shape);
    (*this) = initv;
  }
  TensorContainer
  (const TensorContainer<Device, dimension, DType> &src)
      : pad_(src.pad_) {
    this->dptr_ = data_.dptr_ = NULL;
    this->shape_[0] = 0;
    this->stride_ = 0;
    this->data_.stride_ = 0;
    this->data_.shape_[0] = 0;
    this->stream_ = src.stream_;
    if (src.dptr_ != NULL) {
      this->AllocByShape(src.shape_);
      mshadow::Copy(*this, src, this->stream_);
    }
  }
  ~TensorContainer(void) MSHADOW_THROW_EXCEPTION {
    this->Release();
  }
  inline void Resize(const Shape<dimension> &shape) {
    Shape<2> s2 = shape.FlatTo2D();
    if (s2.shape_[1] > data_.stride_ || s2.shape_[0] > data_.size(0)) {
      this->AllocByShape(shape);
    } else {
      this->shape_ = shape;
      if (this->pad_) {
        this->stride_ = data_.stride_;
      } else {
        this->stride_ = s2.shape_[1];
      }
    }
  }
  inline void Resize(const Shape<dimension> &shape, DType initv) {
    this->Resize(shape);
    (*this) = initv;
  }
  inline void set_pad(bool pad) {
    this->pad_ = pad;
  }
  template<typename TStream>
  inline void SaveBinary(TStream &fo) const { // NOLINT(*)
    mshadow::SaveBinary(fo, *this);
  }
  template<typename TStream>
  inline void LoadBinary(TStream &fi) { // NOLINT(*)
    Tensor<cpu, dimension, DType> tmp;
    mshadow::LoadBinary(fi, &tmp, false);
    this->Resize(tmp.shape_);
    Stream<Device> stream;
    Copy(*this, tmp, &stream);
    mshadow::FreeSpace(&tmp);
  }
  inline TensorContainer &operator=
  (const TensorContainer<Device, dimension, DType> &src) {
    this->pad_ = src.pad_;
    this->stream_ = src.stream_;
    if (src.dptr_ != NULL) {
      this->Resize(src.shape_);
      mshadow::Copy(*this, src, this->stream_);
    }
    return *this;
  }
  inline Tensor<Device, dimension, DType> &operator=(DType s) {
    return this->__assign(s);
  }
  template<typename E>
  inline Tensor<Device, dimension, DType> &
  operator=(const expr::Exp<E, DType, expr::type::kMapper> &exp) {
    return this->__assign(exp);
  }
  template<typename E>
  inline Tensor<Device, dimension, DType> &
  operator=(const expr::Exp<E, DType, expr::type::kChainer> &exp) {
    return this->__assign(exp);
  }
  template<typename E>
  inline Tensor<Device, dimension, DType> &
  operator=(const expr::Exp<E, DType, expr::type::kComplex> &exp) {
    return this->__assign(exp);
  }
  inline void Release(void) {
    if (data_.dptr_ != NULL) {
      this->shape_[0] = 0;
      this->stride_ = 0;
      this->data_.stride_ = 0;
      this->data_.shape_[0] = 0;
      try {
        mshadow::FreeSpace(&data_);
      } catch (const dmlc::Error &e) {
        this->dptr_ = data_.dptr_ = NULL;
        throw e;
      }
      this->dptr_ = data_.dptr_ = NULL;
    }
  }

 private:
  bool pad_;
  Tensor<Device, 2, DType> data_;

  inline void AllocByShape(const Shape<dimension>& shape) {
    if (data_.dptr_ != NULL) this->Release();
    data_.shape_ = shape.FlatTo2D();
    mshadow::AllocSpace(&data_, pad_);
    this->dptr_ = data_.dptr_;
    this->shape_ = shape;
    if (this->pad_) {
      this->stride_ = data_.stride_;
    } else {
      this->stride_ = data_.size(1);
    }
  }
};
}  // namespace mshadow
#endif  // MSHADOW_TENSOR_CONTAINER_H_