io.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#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_