MegEngine基本概念

MegEngine基本概念

基本概念

MegEngine 是基于计算图的深度神经网络学习框架。 本文内容会简要介绍计算图及其相关基本概念，以及在 MegEngine 中的实现。

计算图

结合一个简单的数学表达式来介绍计算图中的基本概念。下图是 y = (w * x) + b 这一数学表达式的计算图表示：

从中可以看到，计算图中存在：

• 数据节点（图中的实心圈）：如输入数据 x 、 w 、b ，运算得到的数据 p ，以及最终的运算输出 y ；
• 计算节点（图中的空心圈）：图中 * 和 + 分别表示计算节点 乘法 和 加法，是施加在数据节点上的运算；
• 边（图中的箭头）：表示数据的流向，体现了数据节点和计算节点之间的依赖关系；

如上便是一个简单的计算图示例。

张量（Tensor）

与 PyTorch，TensorFlow 等深度学习框架类似，MegEngine 使用张量（Tensor）来表示计算图中的数据。

张量（Tensor）可以看做 NumPy 中的数组，它可以是一个标量、向量、矩阵或者多维数组。 可以通过 NumPy 或者 Python List 来创建一个 Tensor 。

执行下列代码并查看相关结果：

In [1]:

import numpy as np

import megengine as mge

# 初始化一个维度为 (2, 5) 的 ndarray，并转化成 MegEngine 的 Tensor

# 注：目前 MegEngine Tensor 不支持 float64 数值类型，所以这里显示指定了 ndarray 的数值类型

a = mge.tensor(np.random.random((2,5)).astype('float32'))

print(a)

# 初始化一个长度为3的列表，并转化成 Tensor

b = mge.tensor([1., 2., 3.])

print(b)

Tensor([[0.2397 0.6569 0.8097 0.3523 0.4676]

[0.9492 0.9415 0.5636 0.614  0.8424]], device=xpux:0)

Tensor([1. 2. 3.], device=xpux:0)

通过 dtype 属性，可以获取 Tensor 的数据类型；通过 astype() 方法，可以拷贝创建一个指定数据类型的新 Tensor ，原 Tensor 不变。

In [2]:

print(a.dtype)

d = a.astype("float16")

print(d.dtype)

<class 'numpy.float32'>

<class 'numpy.float16'>

通过 shape 属性，可以获取 Tensor 的形状：

In [3]:

print(a.shape)

(2, 5)

通过 numpy() 方法，可以将 Tensor 转换为 numpy.ndarray：

In [4]:

a = mge.tensor(np.random.random((2,5)).astype('float32'))

print(a)

b = a.numpy()

print(b)

Tensor([[0.8246 0.8447 0.3225 0.2583 0.6065]

[0.4701 0.594  0.1612 0.7749 0.0067]], device=xpux:0)

[[0.8246328  0.8447421  0.32254046 0.25825405 0.60646415]

[0.47006    0.59400123 0.16122891 0.77490866 0.0067136 ]]

通过 device 属性，可以查询当前 Tensor 所在的设备。创建的 Tensor 可以位于不同 device，这根据当前的环境决定。一般地，如果在创建 Tensor 时不指定 device，其 device 属性默认为 xpux，表示当前任意一个可用的设备。如果存在 GPU 则优先使用 GPU，否则为 CPU。

In [5]:

print(a.device)

xpux:0

也可以在创建 Tensor 时，指定 device 为 cpu0, cpu1, …, gpu0, gpu1, … ，也可以是 cpux 或 gpux，表示当前任意一个可用的 CPU 或 GPU。

通过 to() 方法可以在另一个 device 上生成当前 Tensor 的拷贝，比如将刚刚创建的 Tensor a 迁移到 CPU 上，再迁移到 GPU 上：

In [6]:

b = a.to("cpu0")

print(b.device)

cpu0:0

GPU 和 CPU 切换

MegEngine 在 GPU 和 CPU 同时存在时默认使用 GPU 进行训练。用户可以调用 set_default_device() 来根据自身需求设置默认计算设备。

如下代码设置默认设备为 CPU：

In [7]:

import megengine as mge

# 默认使用 CPU

mge.set_default_device('cpux')

如下代码设置默认设备为GPU:

mge.set_default_device('gpux')

如果不想修改代码，用户也可通过环境变量 MGE_DEFAULT_DEVICE 来设置默认计算设备：

export MGE_DEFAULT_DEVICE='cpux'

export MGE_DEFAULT_DEVICE='gpux'

算子（Operator）

MegEngine 中通过算子 (Operator） 来表示运算。 类似于 NumPy，MegEngine 中的算子支持基于 Tensor 的常见数学运算和操作。 下面介绍几个简单的示例：

Tensor 的加法：

In [8]:

a = mge.tensor(np.random.random((2,5)).astype('float32'))

print(a)

b = mge.tensor(np.random.random((2,5)).astype('float32'))

print(b)

print(a + b)

Tensor([[0.1337 0.5079 0.0929 0.0834 0.7817]

[0.461  0.7906 0.81   0.6579 0.5813]], device=cpux:0)

Tensor([[0.8897 0.451  0.6765 0.3549 0.88  ]

[0.4421 0.7505 0.6881 0.9912 0.6448]], device=cpux:0)

Tensor([[1.0234 0.9589 0.7694 0.4383 1.6617]

[0.9031 1.5411 1.4981 1.6491 1.2261]], device=cpux:0)

Tensor 的切片：

In [9]:

print(a[1, :])

Tensor([0.461  0.7906 0.81   0.6579 0.5813], device=cpux:0)

Tensor 形状的更改：

In [10]:

a.reshape(5, 2)

Out[10]:

Tensor([[0.1337 0.5079]

[0.0929 0.0834]

[0.7817 0.461 ]

[0.7906 0.81  ]

[0.6579 0.5813]], device=cpux:0)

reshape() 的参数允许存在单个缺省值，用 -1 表示。此时，reshape 会自动推理该维度的值：

In [11]:

# 原始维度是 (2, 5)，当给出 -1的缺省维度值时，可以推理出另一维度为10

a = a.reshape(1, -1)

print(a.shape)

(1, 10)

创建的Tensor可以位于不同device，这根据当前的环境决定。通过 device 属性查询当前 Tensor 所在的设备。

In [12]:

print(a.device)

cpux:0

以上是一些简单操作。 可以在 MegEngine API 文档中查询更多算子的用法，比如矩阵乘，卷积等。

求导器（Grad Manager）

神经网络的优化通常通过随机梯度下降来进行。需要根据计算图的输出，对所有的中间数据节点求梯度，这一过程被称之为 “反向传播”，也就是链式求导法则。 例如，希望得到图1中输出 $ y $ 关于输入 $ w $ 的梯度，那么反向传播的过程如下图所示：

首先 $ y = p + b $ ，因此 $ \partial y / \partial p = 1 $ ； 接着，$ p = w * x $ ，因此，$ \partial p / \partial w = x $ 。 根据链式求导法则，$ \partial y / \partial w = \partial y / \partial p * \partial p / \partial w $ ， 因此， $ y $ 关于输入 $ w $ 的梯度为 $ 1 * x = x $ 。

MegEngine 为计算图中的张量提供了自动求导功能，以上图的例子说明： 假设图中的 $ x $ 是 shape 为 (1, 3) 的张量， $ w $ 是 shape 为 (3, 1) 的张量， $ b $ 是一个标量。 利用MegEngine 计算 $ y = x * w + b $ 的过程如下：

In [13]:

import megengine as mge

import megengine.functional as F

from megengine.autodiff import GradManager

x = mge.tensor([1., 3., 5.]).reshape(1, 3)

w = mge.tensor([2., 4., 6.]).reshape(3, 1)

b = mge.tensor(-1.)

gm = GradManager().attach([w, b])   # 新建一个求导器，绑定需要求导的变量

with gm:                            # 开始记录计算图

p = F.matmul(x, w)

y = p + b

gm.backward(y)                  # 计算 y 的导数

print(w.grad)

print(b.grad)

Tensor([[1.]

[3.]

[5.]], device=cpux:0)

Tensor([1.], device=cpux:0)

可以看到，求出的梯度本身也是 Tensor。

with 代码段中的前向运算都会被求导器记录。也可以用 record() 和 release() 来替代 with，分别控制求导器的开启和关闭，代码如下所示。

gm = GradManager().attach([w, b])   # 新建一个求导器，绑定需要求导的变量

gm.record()                         # 开始记录计算图

p = F.matmul(x, w)

y = p + b

gm.backward(y)                      # 计算 y 的导数

gm.release()                        # 停止记录计算图并释放资源

此外，可以使用 detach 方法，把 Tensor 当作一个常量，这样求导器将不会对其求导。如下所示：

In [14]:

x = mge.tensor([1., 3., 5.]).reshape(1, 3)

w = mge.tensor([2., 4., 6.]).reshape(3, 1)

b = mge.tensor(-1.)

gm = GradManager().attach([w, b])   # 新建一个求导器，绑定需要求导的变量

with gm:                            # 开始记录计算图

p = F.matmul(x, w)

y = p + b.detach()              # 停止对 b 求导

gm.backward(y)                  # 计算 y 的导数

print(b.grad)

None

可以看到，梯度本身也是 Tensor。 注意：这里 F.grad() 的第一个参数 y 是个标量，目前尚不支持其为一个多维 Tensor。