统一 MXNet & PyTorch & TensorFlow

2021-02-21 本文已影响0人水之心

为了在 MXNet、PyTorch 以及 TensorFlow 之间进行互转，本文研究它们之间的基础运算的异同。

基础函数

对于 MXNet 有 np（numpy）模块和 npx（numpy_extension）模块。np 模块包含了 NumPy 支持的函数。而 npx 模块包含了一组扩展函数，用来在类似 NumPy 的环境中实现深度学习开发。当使用张量时，我们几乎总是会调用 set_np 函数：这是为了兼容 MXNet 其他的张量处理组件。

from mxnet import np, npx

npx.set_np()

同样，我们可以将 PyTorch 和 TensorFlow 写作 np，这样有：

def import_np(module_name):
    if module_name == 'mxnet':
        from mxnet import np, npx
        npx.set_np()
        np.randn = np.random.randn
        return np
    elif module_name == 'torch':
        import torch as np
        np.array = np.tensor
        np.concatenate = np.cat
        return np
    elif module_name == 'tensorflow':
        from tensorflow.experimental import numpy as np
        return np

这样 MXNet，TensorFlow 与 PyTorch 有相同的函数：