[TensorV2] Add utility member functions to TensorV2 class

This pull request adds several new utility member functions to the TensorV2 class, enabling users to perform various tasks with their tensors more easily and efficiently.
These include saving and loading tensors, updating batches, getting argmax and max absolute values, and more.
The implementation is based on the current Tensor class and aims to improve the overall usability and flexibility of the TensorV2 class.

**Changes proposed in this PR:**
- Added save() and read() methods to allow saving and loading of saved tensor data.
- Added Map() method to create a new Tensor object from a buffer.
- Added argmax() and max_abs() methods to retrieve the indices of max value by batch and the value of the maximum absolute element in a tensor.
- Added updateBatch() to update tensor batch size.

**Self-evaluation:**
1. Build test: [X]Passed [ ]Failed [ ]Skipped
2. Run test:   [X]Passed [ ]Failed [ ]Skipped

Signed-off-by: Donghyeon Jeong <[email protected]>

nntrainer/tensor/float_tensor.cpp

@@ -790,6 +790,28 @@ void FloatTensor::copyData(const TensorV2 &from) {
   }
 }
 
+std::vector<unsigned int> FloatTensor::argmax() const {
+  std::vector<unsigned int> result;
+  const float *data = (float *)getData();
+  size_t batch_size = batch();
+  size_t feature_len = dim.getFeatureLen();
+
+  result.resize(batch_size);
+
+  for (unsigned int b = 0; b < batch_size; b++) {
+    auto max_iter =
+      std::max_element(data + b * feature_len, data + (b + 1) * feature_len);
+    result[b] = std::distance(data, max_iter) - (b * feature_len);
+  }
+  return result;
+}
+
+float FloatTensor::max_abs() const {
+  const float *data = (float *)getData();
+  unsigned int idx = isamax(size(), data, 1);
+  return *(data + idx);
+}
+
 TensorV2 &FloatTensor::transpose(const std::string &direction,
                                  TensorV2 &output) const {
   unsigned int SL, SI, SJ, SK;

nntrainer/tensor/float_tensor.h

@@ -338,6 +338,16 @@ class FloatTensor : public TensorBase {
    */
   void copyData(const TensorV2 &from);
 
+  /**
+   * @copydoc TensorV2::argmax()
+   */
+  std::vector<unsigned int> argmax() const override;
+
+  /**
+   * @copydoc TensorV2::max_abs()
+   */
+  float max_abs() const override;
+
   /**
    * @copydoc TensorV2::transpose(const std::string &direction, TensorV2 &out)
    */

nntrainer/tensor/half_tensor.cpp

@@ -1048,6 +1048,29 @@ void HalfTensor::copyData(const TensorV2 &from) {
   }
 }
 
+std::vector<unsigned int> HalfTensor::argmax() const {
+  std::vector<unsigned int> result;
+  const _FP16 *data = (_FP16 *)getData();
+  size_t batch_size = batch();
+  size_t feature_len = dim.getFeatureLen();
+
+  result.resize(batch_size);
+
+  for (unsigned int b = 0; b < batch_size; b++) {
+    auto max_iter =
+      std::max_element(data + b * feature_len, data + (b + 1) * feature_len);
+    result[b] = std::distance(data, max_iter) - (b * feature_len);
+  }
+
+  return result;
+}
+
+float HalfTensor::max_abs() const {
+  const _FP16 *data = (_FP16 *)getData();
+  unsigned int idx = isamax(size(), data, 1);
+  return (float)(*(data + idx));
+}
+
 TensorV2 &HalfTensor::transpose(const std::string &direction,
                                 TensorV2 &output) const {
   unsigned int SL, SI, SJ, SK;

nntrainer/tensor/half_tensor.h

@@ -337,6 +337,16 @@ class HalfTensor : public TensorBase {
    */
   void copyData(const TensorV2 &from);
 
+  /**
+   * @copydoc TensorV2::argmax()
+   */
+  std::vector<unsigned int> argmax() const override;
+
+  /**
+   * @copydoc TensorV2::max_abs()
+   */
+  float max_abs() const override;
+
   /**
    * @copydoc TensorV2::transpose(const std::string &direction, TensorV2 &out)
    */

nntrainer/tensor/tensor_base.cpp

@@ -40,13 +40,28 @@ bool TensorBase::operator==(const TensorBase &rhs) const {
   return true;
 }
 
+void TensorBase::setTensorVar(TensorDim d, void *buf, size_t offset) {
+  dim = d;
+  strides = d.computeStrides();
+  /// Tensor does not own the memory
+  data = std::shared_ptr<MemoryData>(new MemoryData((void *)buf),
+                                     std::default_delete<MemoryData>());
+  offset = offset;
+}
+
 void TensorBase::putData() const {
   if (!data)
     return;
 
   data->invalidate();
 }
 
+const std::shared_ptr<MemoryData> TensorBase::getMemoryData() const {
+  return data;
+}
+
+size_t TensorBase::getOffset() const { return offset; }
+
 void TensorBase::reshape(const TensorDim &d) {
   NNTR_THROW_IF(!contiguous, std::invalid_argument)
     << getName() << " is not contiguous, cannot reshape.";
@@ -64,6 +79,16 @@ void TensorBase::reshape(const TensorDim &d) {
   strides = d.computeStrides();
 }
 
+void TensorBase::updateBatch(unsigned int batch) {
+  if (dim.batch() == batch) {
+    return;
+  }
+
+  if (isAllocated())
+    throw std::invalid_argument("Cannot update batch for an allocated tensor");
+  dim.batch(batch);
+}
+
 size_t TensorBase::getIndex(unsigned int b, unsigned int c, unsigned int h,
                             unsigned int w) const noexcept {
   if (getFormat() == Tformat::NCHW) {

nntrainer/tensor/tensor_base.h

@@ -128,6 +128,11 @@ class TensorBase {
    */
   bool operator!=(const TensorBase &rhs) const { return !(*this == rhs); }
 
+  /**
+   * @copydoc TensorV2::setTensorVar(TensorDim d, void *buf, size_t offset)
+   */
+  void setTensorVar(TensorDim d, void *buf, size_t offset);
+
   /**
    * @brief Basic Destructor
    */
@@ -369,6 +374,16 @@ class TensorBase {
    */
   virtual void copyData(const TensorV2 &from) = 0;
 
+  /**
+   * @copydoc TensorV2::argmax()
+   */
+  virtual std::vector<unsigned int> argmax() const = 0;
+
+  /**
+   * @copydoc TensorV2::max_abs()
+   */
+  virtual float max_abs() const = 0;
+
   /**
    * @copydoc TensorV2::transpose(const std::string &direction, TensorV2 &out)
    */
@@ -381,6 +396,17 @@ class TensorBase {
    */
   void putData() const;
 
+  /**
+   * @brief     return Data pointer of Tensor
+   * @retval    template T pointer (float pointer as default)
+   */
+  const std::shared_ptr<MemoryData> getMemoryData() const;
+
+  /**
+   * @brief     return offset
+   */
+  size_t getOffset() const;
+
   /**
    * @brief     set Tensor Dim
    * @param[in] d TensorDim
@@ -417,6 +443,12 @@ class TensorBase {
    */
   Tdatatype getDataType() const { return dim.getDataType(); }
 
+  /**
+   * @brief     update batch size for this tensor
+   * @param     batch size
+   */
+  void updateBatch(unsigned int batch);
+
   /**
    * @brief     return whether tensor is contiguous or not.
    * @retval    bool contiguous

nntrainer/tensor/tensor_v2.cpp

@@ -736,6 +736,12 @@ void TensorV2::print(std::ostream &out) const { itensor->print(out); }
 
 void TensorV2::putData() const { itensor->putData(); }
 
+const std::shared_ptr<MemoryData> TensorV2::getMemoryData() const {
+  return itensor->getMemoryData();
+}
+
+size_t TensorV2::getOffset() const { return itensor->getOffset(); }
+
 void TensorV2::copy(const TensorV2 &from) {
   /// @todo enable copy to non-contiguous tensor
   if (!itensor->getContiguous()) {
@@ -807,6 +813,45 @@ TensorV2 TensorV2::clone() const {
   return output;
 }
 
+void TensorV2::save(std::ostream &file) {
+  NNTR_THROW_IF(!getContiguous(), std::invalid_argument)
+    << getName() << " is not contiguous, cannot save.";
+
+  std::streamsize sz = static_cast<std::streamsize>(bytes());
+  NNTR_THROW_IF(sz < 0, std::invalid_argument)
+    << "save size: " << bytes()
+    << " is too big. It cannot be represented by std::streamsize";
+
+  checkedWrite(file, getData<char>(), sz, "[Tensor::save] operation failed");
+  putData();
+}
+
+void TensorV2::read(std::ifstream &file) {
+  NNTR_THROW_IF(!getContiguous(), std::invalid_argument)
+    << getName() << " is not contiguous, cannot read.";
+
+  std::streamsize sz = static_cast<std::streamsize>(bytes());
+
+  NNTR_THROW_IF(sz < 0, std::invalid_argument)
+    << "read size: " << bytes()
+    << " is too big. It cannot be represented by std::streamsize";
+
+  checkedRead(file, getData<char>(), sz, "[Tensor::read] operation failed");
+  putData();
+}
+
+std::vector<unsigned int> TensorV2::argmax() const {
+  NNTR_THROW_IF(!getContiguous(), std::invalid_argument)
+    << getName() << " is not contiguous, cannot get argmax.";
+  return itensor->argmax();
+}
+
+float TensorV2::max_abs() const {
+  NNTR_THROW_IF(!getContiguous(), std::invalid_argument)
+    << getName() << " is not contiguous, cannot get max_abs.";
+  return itensor->max_abs();
+}
+
 TensorV2 TensorV2::transpose(const std::string &direction) const {
   TensorV2 output(getDim());
   transpose(direction, output);
@@ -844,6 +889,8 @@ TensorDim::Format TensorV2::getFormat() const { return itensor->getFormat(); }
 
 Tdatatype TensorV2::getDataType() const { return itensor->getDataType(); }
 
+void TensorV2::updateBatch(unsigned int batch) { itensor->updateBatch(batch); }
+
 const bool TensorV2::getContiguous() const noexcept {
   return itensor->getContiguous();
 }
@@ -920,4 +967,13 @@ TensorV2 TensorV2::getSharedDataTensor(const TensorDim dim_, size_t offset,
   return ret;
 }
 
+void TensorV2::setTensorVar(TensorDim d, void *buf, size_t offset) {
+  itensor->setTensorVar(d, buf, offset);
+}
+
+std::ostream &operator<<(std::ostream &out, TensorV2 const &input) {
+  input.print(out);
+  return out;
+}
+
 } // namespace nntrainer