本文首发于 MXNet与Gluon的黑科技工具箱写文章
作者:梁德澎

前言

本文主要内容是解读MXNet是加载模型参数文件并解析得到NDArray所涉及到的代码,希望读者读完本文能对MXNet参数文件的格式有清晰的了解,并可以自己来实现参数文件的解析。

解析MXNet参数文件C++小工程:

https://github.com/Ldpe2G/DeepLearningForFun/tree/master/MXNet-Cpp/parsingNDArray​github.com

本文解读的MXNet代码基于的版本:commit 7d2c9bf3b631433132452760734b684e39170814

Python前端代码入口

首先从MXNet Python前端看是如何是调用底层C接口来读取NDArray参数文件的,这部分代码见源码 ${MXNET_ROOT}/python/mxnet/ndarray/utils.py 第149行:

def load(fname):
    if not isinstance(fname, string_types):
        raise TypeError('fname required to be a string')
    out_size = mx_uint()
    out_name_size = mx_uint()
    handles = ctypes.POINTER(NDArrayHandle)()
    names = ctypes.POINTER(ctypes.c_char_p)()
    check_call(_LIB.MXNDArrayLoad(c_str(fname),
                                  ctypes.byref(out_size),
                                  ctypes.byref(handles),
                                  ctypes.byref(out_name_size),
                                  ctypes.byref(names)))
    .....

这个 load 函数接收参数路径作为输入,主要是调用 MXNDArrayLoad 这个底层的C函数接口来读取参数。MXNet底层实现是C++,并提供了一层C的接口供前端语言去调用。

C接口层

接着来看下MXNDArrayLoad接口的实现,这部分代码见${MXNET_ROOT}/src/c_api/c_api.cc第1344行:

int MXNDArrayLoad(const char* fname,
                  uint32_t *out_size,
                  NDArrayHandle** out_arr,
                  uint32_t *out_name_size,
                  const char*** out_names) {
  MXAPIThreadLocalEntry<> *ret = MXAPIThreadLocalStore<>::Get();
  ret->ret_vec_str.clear();
  API_BEGIN();
  std::vector<NDArray> data;
  std::vector<std::string> &names = ret->ret_vec_str;
  {
    std::unique_ptr<dmlc::Stream> fi(dmlc::Stream::Create(fname, "r"));
    mxnet::NDArray::Load(fi.get(), &data, &names);
  }
  ......
}

核心代码就是首先打开文件流,接着调用NDArray类的静态函数mxnet::NDArray::Load函数读取并解析参数文件,得到参数NDArray数组保存到data这个变量里面。我们只要关注dmlc::Stream这个类的实现还有mxnet::NDArray::Load这个类的实现就可以了。

底层C++实现

NDArray::Load静态函数

函数具体实现见${MXNET_ROOT}/src/ndarray/ndarray.cc第 1924 行:

void NDArray::Load(dmlc::Stream* fi,
                   std::vector<NDArray>* data,
                   std::vector<std::string>* keys) {
  uint64_t header, reserved;
  CHECK(fi->Read(&header))
      << "Invalid NDArray file format";
  CHECK(fi->Read(&reserved))
      << "Invalid NDArray file format";
  CHECK(header == kMXAPINDArrayListMagic)
      << "Invalid NDArray file format";
  CHECK(fi->Read(data))
      << "Invalid NDArray file format";
  CHECK(fi->Read(keys))
      << "Invalid NDArray file format";
  CHECK(keys->size() == 0 || keys->size() == data->size())
      << "Invalid NDArray file format";
}

从这里读取内容的过程可以大概看出NDArray参数文件存储的内容格式。

首先文件开头保存了两个uint64\_t类型的数字,接着就是NDArray参数数组,接着是每个NDArray对应的名字数组。

读取的时候都是调用Strem类的Read函数,接下来就是看下Stream类的实现。

Stream类

看回上面打开参数文件的代码:

std::unique_ptr<dmlc::Stream> fi(dmlc::Stream::Create(fname, "r"));

dmlc::Stream::Create代码见dmlc-core子模块:${MXNET_ROOT}/3rdparty/dmlc-core/src/io.cc第132行:

Stream *Stream::Create(const char *uri,
                       const char * const flag,
                       bool try_create) {
  io::URI path(uri);
  return io::FileSystem::
      GetInstance(path)->Open(path, flag, try_create);
}

调用了FileSystem::GetINstance函数得到实例,并调用Open函数打开文件,这里返回的实例是LocalFileSystem类的实例,其Open函数见${MXNET_ROOT}/3rdparty/dmlc-core/src/io/local_filesys.cc第147行:

SeekStream *LocalFileSystem::Open(const URI &path,
                                  const char* const mode,
                                  bool allow_null) {
  FILE *fp = NULL;
  const char *fname = path.name.c_str();
  using namespace std;
  std::string flag = mode;
  if (flag == "r") flag = "rb";
  fp = fopen(fname, flag.c_str());
  if (fp != NULL) {
    return new FileStream(fp, false);
  } else {
    return NULL;
  }
}

为了可读性我简化了代码,可以看到就是调用std::fopen函数打开文件,并把FILE指针传给FileStream类,

代码见${MXNET_ROOT}/3rdparty/dmlc-core/src/io/local_filesys.cc第27行:

class FileStream : public SeekStream {
 public:
  explicit FileStream(FILE *fp, bool use_stdio)
      : fp_(fp), use_stdio_(use_stdio) {}
  virtual ~FileStream(void) {
    this->Close();
  }
  virtual size_t Read(void *ptr, size_t size) {
    return std::fread(ptr, 1, size, fp_);
  }

  ......
  
 private:
  std::FILE *fp_;
  bool use_stdio_;
};

可以看到FileStream继承自SeekStrem,而且成员函数Read实现的功能是调用std::fread函数从fp_文件指针里面读取size大小字节的内容,std::fread的文档见https://en.cppreference.com/w/cpp/io/c/fread

看下每个参数的解释:

buffer - void 指针指向从文件留中读取内容的存取内存
  size   - 指针指向内存每个元素字节大小,这里由于是void指针,所以size大小恒为1
  count  - 读取元素个数
  stream - 文件流

MXNet这里的实现是把需要被读取的内存指针转换成void *,这样子就可以兼容各种基本类型的指针读取,只需要记住传入的读取元素个数是 sizeof(T) * count,就是原来类型元素个数乘以每个元素对应的字节数。

接着再看回上面打开参数文件的代码:

std::unique_ptr<dmlc::Stream> fi(dmlc::Stream::Create(fname, "r"));

Strem::Create返回的是Strem类型,而不是SeekStrem,所以继续往上找Strem类的定义,代码见${MXNET_ROOT}/3rdparty/dmlc-core/include/dmlc/io.h第30行:

class Stream {  // NOLINT(*)
 public:

  virtual size_t Read(void *ptr, size_t size) = 0;

  static Stream *Create(const char *uri,
                        const char* const flag,
                        bool allow_null = false);
  
  template<typename T>
  inline bool Read(T *out_data);

  template<typename T>
  inline bool ReadArray(T* data, size_t num_elems);
};

为了可读性,只保留了读文件相关的代码,可以看到,FileStream重写了virtual size_t Read(void *ptr, size_t size) = 0;虚函数,

而回看NDArray静态Load函数:

void NDArray::Load(dmlc::Stream* fi,
                   std::vector<NDArray>* data,
                   std::vector<std::string>* keys) {
  uint64_t header, reserved;
  CHECK(fi->Read(&header))
      << "Invalid NDArray file format";
  CHECK(fi->Read(&reserved))
      << "Invalid NDArray file format";
  CHECK(header == kMXAPINDArrayListMagic)
      << "Invalid NDArray file format";
  CHECK(fi->Read(data))
      << "Invalid NDArray file format";
  CHECK(fi->Read(keys))
      << "Invalid NDArray file format";
  CHECK(keys->size() == 0 || keys->size() == data->size())
      << "Invalid NDArray file format";
}

具体读参数文件内容的时候调用的是Stream类的Read(T *out_data)模板函数

  template<typename T>
  inline bool Read(T *out_data);

这个模板函数的实现很有意思,代码见${MXNET_ROOT}/3rdparty/dmlc-core/include/dmlc/io.h第455:

template<typename T>
inline bool Stream::Read(T *out_data) {
  return serializer::Handler<T>::Read(this, out_data);
}

可以看到调用了Handler::Read函数,继续跟进去看Handler的实现,代码见${MXNET_ROOT}/3rdparty/dmlc-core/include/dmlc/serializer.h第258行:

template<typename T>
struct Handler {
  ......
  /*!
   * \brief read data to stream
   * \param strm the stream to read the data.
   * \param data the pointer to the data obeject to read
   * \return whether the read is successful
   */
  inline static bool Read(Stream *strm, T *data) {
    return
    IfThenElse<dmlc::is_arithmetic<T>::value,
               ArithmeticHandler<T>,
               IfThenElse<dmlc::is_pod<T>::value && DMLC_IO_NO_ENDIAN_SWAP,
                          NativePODHandler<T>,
                          IfThenElse<dmlc::has_saveload<T>::value,
                                     SaveLoadClassHandler<T>,
                                     UndefinedSerializerFor<T>, T>,
                          T>,
               T>
    ::Read(strm, data);
  }
};

一开始看到这串代码可能会有点懵,不过没关系,接下来我们就一步步拆解这段代码,首先看IfThenElse结构体的实现,代码见${MXNET_ROOT}/3rdparty/dmlc-core/include/dmlc/serializer.h第48行:

template<bool cond, typename Then, typename Else, typename Return>
struct IfThenElse;

template<typename Then, typename Else, typename T>
struct IfThenElse<true, Then, Else, T> {
  ......
  inline static bool Read(Stream *strm, T *data) {
    return Then::Read(strm, data);
  }
};
template<typename Then, typename Else, typename T>
struct IfThenElse<false, Then, Else, T> {
  ......
  inline static bool Read(Stream *strm, T *data) {
    return Else::Read(strm, data);
  }
};

我理解就是根据模板参数在编译期间做分支选择,根据模板参数决定调用实现分支,这里可以看到如果第一个模板参数template<bool cond, ...>true的话,就调用 Then::Read函数 ,否则调用Else::Read函数,然后看回Handler::Read函数:

  inline static bool Read(Stream *strm, T *data) {
    return
    IfThenElse<dmlc::is_arithmetic<T>::value,
               ArithmeticHandler<T>,
               IfThenElse<dmlc::is_pod<T>::value && DMLC_IO_NO_ENDIAN_SWAP,
                          NativePODHandler<T>,
                          IfThenElse<dmlc::has_saveload<T>::value,
                                     SaveLoadClassHandler<T>,
                                     UndefinedSerializerFor<T>, T>,
                          T>,
               T>
    ::Read(strm, data);
  }
};

代码就很好理解了,如果dmlc::is_arithmetic<T>::value值为true则走ArithmeticHandler<T>,否则再进行第二次判断,先来看下具体实现${MXNET_ROOT}/3rdparty/dmlc-core/include/dmlc/type_traits.h第66行:

template<typename T>
struct is_arithmetic {
  static const bool value = std::is_arithmetic<T>::value;
};

为了可读性我简化了下代码,只要查下C++的文档看下std::is_arithmetic<T>的定义就知道什么模板参数类型是什么的情况下值是true或者false,C++文档解释见https://en.cppreference.com/w/cpp/types/is_arithmetic

也就是如果模板类型T是整数型或者浮点型,value的值就是true否则是false。如果满足条件则会调用ArithmeticHandler::Read函数,代码见${MXNET_ROOT}/3rdparty/dmlc-core/include/dmlc/serializer.h第82行:

/*! \brief Serializer for arithmetic data, handle endianness */
template<typename T>
struct ArithmeticHandler {
  ......
  inline static bool Read(Stream *strm, T *dptr) {
    bool ret = strm->Read((void*)dptr, sizeof(T)) == sizeof(T);  
    ......
    return ret;
  }
};

就是运行时调用子类重写的Read函数,从文件流中读取一个T类型元素。接着再看回其他分支选择:

  inline static bool Read(Stream *strm, T *data) {
    return
    IfThenElse<dmlc::is_arithmetic<T>::value,
               ArithmeticHandler<T>,
               IfThenElse<dmlc::is_pod<T>::value && DMLC_IO_NO_ENDIAN_SWAP,
                          NativePODHandler<T>,
                          IfThenElse<dmlc::has_saveload<T>::value,
                                     SaveLoadClassHandler<T>,
                                     UndefinedSerializerFor<T>, T>,
                          T>,
               T>
    ::Read(strm, data);
  }
};

如果不满足ArithmeticHandler的条件,则看下面一个判断dmlc::is_pod<T>,代码见${MXNET_ROOT}/3rdparty/dmlc-core/include/dmlc/type_traits.h第20行:

template<typename T>
struct is_pod {
#if DMLC_USE_CXX11
  /*! \brief the value of the traits */
  static const bool value = std::is_pod<T>::value;
#else
  /*! \brief the value of the traits */
  static const bool value = false;
#endif
};

看C++文档解释https://en.cppreference.com/w/cpp/types/is_pod

如果是plain old data type值就是true,关于PODType的解释大家可以参考:

https://zhuanlan.zhihu.com/p/29734547

https://en.cppreference.com/w/cpp/named\_req/PODType

NativePODHandler::Read函数的实现也是和ArithmeticHandler::Read类似,也是运行时调用Stream子类重写的Read函数,从文件流中读取一个T类型元素

template<typename T>
struct NativePODHandler {
  ......
  inline static bool Read(Stream *strm, T *dptr) {
    return strm->Read((void*)dptr, sizeof(T)) == sizeof(T);  // NOLINT(*)
  }
};

接着继续回看下一个条件判断dmlc::has_saveload<T>

  inline static bool Read(Stream *strm, T *data) {
    return
    IfThenElse<......
                          IfThenElse<dmlc::has_saveload<T>::value,
                                     SaveLoadClassHandler<T>,
                                     UndefinedSerializerFor<T>, T>,
                          T>,
               T>
    ::Read(strm, data);
  }
};

代码见${MXNET_ROOT}/3rdparty/dmlc-core/include/dmlc/type_traits.h第109行:

template<typename T>
struct has_saveload {
  /*! \brief the value of the traits */
  static const bool value = false;
};

默认值是false,不过看到${MXNET_ROOT}/include/mxnet/ndarray.h第1492行:

namespace dmlc {
/*!\brief traits */
DMLC_DECLARE_TRAITS(has_saveload, mxnet::NDArray, true);
}  // namespace dmlc

${MXNET_ROOT}/3rdparty/dmlc-core/include/dmlc/type_traits.h第125行,DMLC_DECLARE_TRAITS的宏定义:

/*! \brief macro to quickly declare traits information */
#define DMLC_DECLARE_TRAITS(Trait, Type, Value)       \
  template<>                                          \
  struct Trait<Type> {                                \
    static const bool value = Value;                  \
  }

就知道对于NDArray类来说,dmlc::has_saveload<NDArray>::value == true,所以可以判断如果模板参数类型是NDArray则会进入SaveLoadClassHandler实现:

template<typename T>
struct SaveLoadClassHandler {
  ......
  inline static bool Read(Stream *strm, T *data) {
    return data->Load(strm);
  }
};

实际就是调用了T::Load函数,也就是NDArray::Load成员函数。

Handler类还提供了其他模板参数类型的支持比如vector<T>或者std::string

template<typename T>
struct Handler<std::vector<T> > {
  ......
  inline static bool Read(Stream *strm, std::vector<T> *data) {
    return IfThenElse<dmlc::is_pod<T>::value && DMLC_IO_NO_ENDIAN_SWAP,
                      NativePODVectorHandler<T>,
                      ComposeVectorHandler<T>,
                      std::vector<T> >
    ::Read(strm, data);
  }
};

template<typename T>
struct Handler<std::basic_string<T> > {
  .....
  inline static bool Read(Stream *strm, std::basic_string<T> *data) {
    return IfThenElse<dmlc::is_pod<T>::value && (DMLC_IO_NO_ENDIAN_SWAP || sizeof(T) == 1),
                      NativePODStringHandler<T>,
                      UndefinedSerializerFor<T>,
                      std::basic_string<T> >
    ::Read(strm, data);
  }
};

相信有了前面Handler类的解释,理解这两个模板类的实现也就容易多了,这里就不展开了有兴趣的读者可以继续去深入了解,下面回来继续看参数NDArray加载参数逻辑部分的代码解读。

MXNet参数文件解析逻辑

首先给出MXNet参数文件存储内容的格式示意图:

然后根据官方代码的解析逻辑,我自己实现的参数提取代码,为了可读性简化了代码,完整代码见文章开头的github链接:

struct cpu {
  static const int kDevMask = 1 << 0;
};
struct gpu {
  static const int kDevMask = 1 << 1;
};

enum DeviceType {
  kCPU = cpu::kDevMask,
  kGPU = gpu::kDevMask,
  kCPUPinned = 3,
  kCPUShared = 5,
};

static bool Read(std::FILE *fp, void *ptr, size_t size) {
  return std::fread(ptr, 1, size, fp) == size;
} 

int32_t loadNDArrayV2(std::vector<NDArray *>& ndarrays, std::string param_file) {

  std::FILE *fp = fopen(param_file.c_str(), "rb");

  uint64_t header, reserved;
  Read(fp, (void*)(&header), sizeof(uint64_t));
  Read(fp, (void*)(&reserved), sizeof(uint64_t))

  uint64_t nd_size;
  Read(fp, (void*)(&nd_size), sizeof(uint64_t));

  size_t size = static_cast<size_t>(nd_size);
  ndarrays.resize(size);

  // read nd data
  for (size_t i = 0; i < nd_size; ++i) {
    NDArray* nd = new NDArray;
    ndarrays[i] = nd;

    uint32_t magic;
    Read(fp, (void*)(&magic), sizeof(uint32_t));

    // load storage type
    int32_t stype;
    Read(fp, (void*)(&stype), sizeof(int32_t));

    // load shape
    uint32_t ndim_{0};
    Read(fp, (void*)(&ndim_), sizeof(uint32_t));

    size_t nread = sizeof(int64_t) * ndim_;
    int64_t *data_heap_ = new int64_t[ndim_];
    Read(fp, (void*)data_heap_, nread);

    int64_t size = 1;
    for (uint32_t i=0; i<ndim_;++i) {
      size *= data_heap_[i];
      nd->shape.push_back(data_heap_[i]);
    }

    delete[] data_heap_;

    // load context 
    DeviceType dev_type;
    int32_t dev_id;
    Read(fp, (void*)(&dev_type), sizeof(dev_type));

    Read(fp, (void*)(&dev_id), sizeof(int32_t));

    // load type flag
    int32_t type_flag;
    Read(fp, (void*)(&type_flag), sizeof(int32_t));

    size_t all_size = size * mshadow_sizeof(type_flag);
    nd->numOfBytes = all_size;
    nd->data = (void *)malloc(all_size);

    Read(fp, nd->data, nd->numOfBytes);
  }

  // read nd names
  std::vector<std::string> keys;
  uint64_t keysLen;
  Read(fp, (void*)(&keysLen), sizeof(uint64_t));
  keys.resize(keysLen);

  for (uint64_t k = 0; k < keysLen; ++k) {
     uint64_t stringLen;
     Read(fp, (void*)(&stringLen), sizeof(uint64_t));
    size_t size = static_cast<size_t>(stringLen);
    keys[k].resize(size);
    if (size != 0) {
      size_t nbytes = sizeof(char) * size;
      Read(fp, (void*)(&(keys[k][0])), nbytes);
    }
  }
  std::fclose(fp);
  return kSuccess;
}

结合示意图和代码,应该就能比较好的理解参数文件的存储格式和解析方法了。

推荐文章

  • 估算卷积层与反卷积层运算量
  • MXNet实现卷积神经网络训练量化
  • 深入理解深度学习中的反(转置)卷积

更多AI移动端优化的请关注专栏嵌入式AI以及知乎(@梁德澎)。

发表评论

邮箱地址不会被公开。 必填项已用*标注