Tensor.h 1.3 KB
Newer Older
D
dzhwinter 已提交
1
#pragma once
2 3 4 5 6
/**
 * @brief tensor used by optimizer
 */

#include <string.h>
D
dzhwinter 已提交
7
#include <memory>
8
#include "paddle/math/MemoryHandle.h"
D
dzhwinter 已提交
9 10
#include "paddle/utils/Common.h"
#include "paddle/utils/Logging.h"
11 12 13 14 15

namespace paddle {
namespace optimizer {

template <class T>
D
dzhwinter 已提交
16
class TensorT {
17
public:
18 19 20 21 22 23 24
  TensorT(size_t size)
      : TensorT(std::make_shared<CpuMemoryHandle>(size * sizeof(float)), size) {
  }
  TensorT(CpuMemHandlePtr handle, size_t size)
      : height_(1),
        width_(size),
        data_(reinterpret_cast<T*>(handle->getBuf())) {}
D
dzhwinter 已提交
25

D
dzhwinter 已提交
26
  TensorT(T* data, size_t size) : height_(1), width_(size), data_(data) {}
D
dzhwinter 已提交
27 28 29

  TensorT(T* data, size_t h, size_t w) : height_(h), width_(w), data_(data) {}

30
  virtual ~TensorT() {}
D
dzhwinter 已提交
31

32
  T* get_buffer() { return this->data_; }
D
dzhwinter 已提交
33

34
  T& operator[](const size_t idx) {
D
dzhwinter 已提交
35
    CHECK(idx >= 0 && idx < this->width_) << "out of index range";
D
dzhwinter 已提交
36 37
    return data_[idx];
  }
38
  T& operator[](const size_t idx) const {
D
dzhwinter 已提交
39 40
    CHECK(idx >= 0 && idx < this->width_) << "out of index range";
    return data_[idx];
41
  }
42
  // TODO: replace with tensorshape
D
dzhwinter 已提交
43
  size_t size() const { return this->width_ * this->height_; }
D
dzhwinter 已提交
44 45 46 47 48

protected:
  size_t height_;
  size_t width_;
  T* data_;
49 50
};

D
dzhwinter 已提交
51
// TODO(zhihong): design problem of dynamic datatype, need to fix it
52
typedef TensorT<float> Tensor;
D
dzhwinter 已提交
53

54 55
}  // namespace optimizer
}  // namespace paddle