io.h

Go to the documentation of this file.

#ifndef MSHADOW_IO_H_
#define MSHADOW_IO_H_
#include "./tensor.h"

namespace mshadow {
namespace utils {
class IStream {
 public:
  virtual size_t Read(void *ptr, size_t size) = 0;
  virtual void Write(const void *ptr, size_t size) = 0;
  virtual ~IStream(void) {}
};
}  // namespace utils
template<int dim, typename DType, typename TStream>
inline void SaveBinary(TStream &fo, const Tensor<cpu, dim, DType> &src);  // NOLINT(*)
template<int dim, typename DType, typename TStream>
inline void SaveBinary(TStream &fo, const Tensor<gpu, dim, DType> &src); // NOLINT(*)
template<int dim, typename DType, typename TStream>
inline void LoadBinary(TStream &fi,  // NOLINT(*)
                       Tensor<cpu, dim, DType> *dst, bool pre_alloc);
template<int dim, typename DType, typename TStream>
inline void LoadBinary(TStream &fi, // NOLINT(*)
                       Tensor<gpu, dim, DType> *dst, bool pre_alloc);

// implementations
template<int dim, typename DType, typename TStream>
inline void SaveBinary(TStream &fo, const Tensor<cpu, dim, DType> &src_) { // NOLINT(*)
  fo.Write(&src_.shape_, sizeof(src_.shape_));
  Tensor<cpu, 2, DType> src = src_.FlatTo2D();
  for (index_t i = 0; i < src.size(0); ++i) {
    fo.Write(src[i].dptr_, sizeof(DType) * src.size(1));
  }
}
template<int dim, typename DType, typename TStream>
inline void SaveBinary(TStream &fo, const Tensor<gpu, dim, DType> &src) { // NOLINT(*)
  // copy to CPU, then save
  Tensor<cpu, dim, DType> tmp(src.shape_);
  AllocSpace(&tmp);
  Stream<gpu> stream;
  Copy(tmp, src, &stream);
  SaveBinary(fo, tmp);
  FreeSpace(&tmp);
}
template<int dim, typename DType, typename TStream>
inline void LoadBinary(TStream &fi, // NOLINT(*)
                       Tensor<cpu, dim, DType> *dst_, bool pre_alloc) {
  Shape<dim> shape;
  CHECK_NE(fi.Read(&shape, sizeof(shape)), 0) << "mshadow::LoadBinary";
  if (pre_alloc) {
    CHECK_EQ(shape, dst_->shape_) << "LoadBinary, shape do not match pre-allocated shape";
  } else {
    dst_->shape_ = shape; AllocSpace(dst_);
  }
  Tensor<cpu, 2, DType> dst = dst_->FlatTo2D();
  if (dst.size(0) == 0) return;
  for (index_t i = 0; i < dst.size(0); ++i) {
    CHECK_NE(fi.Read(dst[i].dptr_, sizeof(DType) * dst.size(1)), 0) << "mshadow::LoadBinary";
  }
}
template<int dim, typename DType, typename TStream>
inline void LoadBinary(TStream &fi, // NOLINT(*)
                       Tensor<gpu, dim, DType> *dst, bool pre_alloc) {
  Tensor<cpu, dim, DType> tmp;
  LoadBinary(fi, &tmp, false);
  if (pre_alloc) {
    CHECK_EQ(tmp.shape, dst->shape_) << "LoadBinary, shape do not match pre-allocated shape";
  } else {
    dst->shape = tmp.shape; AllocSpace(dst);
  }
  Stream<gpu> stream;
  Copy(*dst, tmp, &stream);
  FreeSpace(&tmp);
}
}  // namespace mshadow
#endif  // MSHADOW_IO_H_