基本概念

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

计算图(Computational Graph)

下面通过一个简单的数学表达式 \(y = (w * x) + b\) 来介绍计算图的基本概念,如下图所示:

../_images/computer_graph.png

图1

从中可以看到,计算图中存在:

  • 数据节点(图中的实心圈):如输入数据 \(x\)\(w\)\(b\) ,运算得到的中间数据 \(p\) ,以及最终的运算输出 \(y\)

  • 计算节点(图中的空心圈):图中 * 和 + 分别表示计算节点 乘法加法,是施加在数据节点上的运算;

  • 边(图中的箭头):表示数据的流向,体现了数据节点和计算节点之间的依赖关系;

如上,便是一个简单的计算图示例。计算图是一个包含数据节点和计算节点的有向图(可以是有环的,也可以是无环的), 是数学表达式的形象化表示。在深度学习领域,任何复杂的深度神经网络本质上都可以用一个计算图表示出来。

前向传播 是计算由计算图表示的数学表达式的值的过程。在图1中,变量 \(x\)\(w\) ,从左侧输入,首先经过乘法运算得到中间结果 \(p\) , 接着,\(p\) 和输入变量 \(b\) 经过加法运算,得到右侧最终的输出 \(y\) ,这就是一个完整的前向传播过程。

在 MegEngine 中,用张量(Tensor)表示计算图中的数据节点,用算子(Operator)实现数据节点之间的运算。

张量(Tensor)

与 PyTorch,TensorFlow 等深度学习框架类似,MegEngine 使用张量(Tensor)来表示计算图中的数据。 张量(Tensor)可以看做 NumPy 中的数组,它可以是标量、向量、矩阵或者多维数组。 可以通过 NumPy 或者 Python List 来创建一个 Tensor 。

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.2976 0.4078 0.5957 0.3945 0.9413]
[0.7519 0.3313 0.0913 0.3345 0.3256]], device=xpux:0)

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

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

print(a.dtype)
d = a.astype("float16")
print(d.dtype)

输出:

<class 'numpy.float32'>
<class 'numpy.float16'>

通过 shape 属性,我们可以获取 Tensor 的形状:

print(a.shape)

输出为一个Tuple:

(2, 5)

通过 numpy() 方法,我们可以将 Tensor 转换为 numpy.ndarray:

a = mge.tensor(np.arange(12)).reshape(2, 6).astype("float32")
print(a)

b = a.numpy()
print(b)

输出:

Tensor([[ 0.  1.  2.  3.  4.  5.]
[ 6.  7.  8.  9. 10. 11.]], device=xpux:0)

[[ 0.  1.  2.  3.  4.  5.]
[ 6.  7.  8.  9. 10. 11.]]

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

print(a.device)

输出:

xpux:0

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

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

# 下面代码是否能正确执行取决于你当前所在的环境
b = a.to("cpu0")
print(b.device)

c = b.to("gpu0")
print(c.device)

输出:

cpu0:0
gpu0:0

GPU 和 CPU 切换

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

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

import megengine as mge

# 默认使用 CPU
mge.set_default_device('cpux')

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

# 默认使用 GPU
mge.set_default_device('gpux')

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

# 默认使用 CPU
export MGE_DEFAULT_DEVICE='cpux'

# 默认使用 GPU
export MGE_DEFAULT_DEVICE='gpux'

算子(Operator)

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

Tensor 的加法:

a = mge.tensor([[1., 2., 2.], [5., 1., 8.]])
print(a)

b = mge.tensor([[1., 9., 1.], [1., 7., 9.]])
print(b)

print(a + b)

输出:

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

Tensor([[1. 9. 1.]
[1. 7. 9.]], device=xpux:0)

Tensor([[ 2. 11.  3.]
[ 6.  8. 17.]], device=xpux:0)

Tensor 的切片:

print(a[1, :])

输出:

Tensor([5. 1. 8.], device=xpux:0)

Tensor 形状的更改:

a.reshape(3, 2)

输出:

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

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

# 原始维度是 (2, 3),当给出 -1 的缺省维度值时,可以推理出另一维度为 6
a = a.reshape(1, -1)
print(a.shape)

输出:

(1, 6)

MegEngine 的 functional 提供了更多的算子,比如深度学习中常用的矩阵乘操作、卷积操作等。

Tensor 的矩阵乘:

import megengine as mge
import megengine.functional as F

a = mge.tensor(np.arange(6).reshape(2, 3)).astype('float32')
print(a)
b = mge.tensor(np.arange(6, 12).reshape(3, 2)).astype('float32')
print(b)
c = F.matmul(a, b)
print(c)

输出:

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

Tensor([[ 6.  7.]
[ 8.  9.]
[10. 11.]], device=xpux:0)

Tensor([[ 28.  31.]
[100. 112.]], device=xpux:0)

更多算子可以参见 functional 部分的文档。

求导器(Grad Manager)

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

../_images/back_prop.png

图2

首先 \(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\) 的梯度为 \(x\)

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

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=xpux:0)

Tensor([1.], device=xpux: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 当作一个常量,这样求导器将不会对其求导。如下所示:

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