概况
天元（英文名：MegEngine）是旷视科技 3 月 25 日开源的深度学习框架，这一名字取自围棋棋盘中心点的名称，也有向 AlphaGo 致敬之意。听了旷视大佬们在发布会上的介绍，MegEngine 从 2014 年开始研发，作为旷视内部全员使用的框架，MegEngine 是驱动旷视在深度学习领域取得一系列成绩的核心动力。
趁没开学在家，对照官方文档快速上手了一下，这篇文章会记录一下对 MegEngine 初步的一些感受。
- 天元 MegEngine 官方网站：https://megengine.org.cn/
- GitHub repo：https://github.com/MegEngine/MegEngine
- 中文社区：https://discuss.megengine.org.cn/
- MegStudio：https://studio.brainpp.com/
安装
目前 MegEngine 的支持平台还比较单一，仅支持 Linux 环境下安装。对 Windows 用户，官方提示了可以使用 WSL 来运行 MegEngine，不过只支持 CPU 后端。
详细的环境要求是：
- 64 位、16.04 及以上版本的 Ubuntu
- Python 3.5+
- NVIDIA 驱动版本 418.x
旷视 MegStudio 平台提供了 MegEngine 0.3.1 + Python 3.8 的环境，可以很方便地开始使用 MegEngine。
除了 MegStudio，实测 Google Colab 平台也可以成功安装 MegEngine。
MegEngine 安装包本身集成了 CUDA 环境，因此不区分 CPU 和 GPU 版本。
通过 pip 安装 MegEngine：
pip3 install megengine -f https://megengine.org.cn/whl/mge.html
整个安装包 700MB+，下载的速度非常快，服务器在国内果然无惧速度问题。
import 一下来测试安装是否成功，官方对 MegEngine 的习惯性缩写是 mge：
import megengine as mge
对于要参与到 MegEngine 开发贡献的开发者，或者需要使用未进入 release 的功能，则需要从源码安装，可以参考 README 中的指引。
体验
MegStudio
MegStudio 是旷视开放的提供免费算力的在线深度学习开发平台。目前提供下面三种配置的环境，其中基础版的环境不限时长，高级版（CPU）和高级版（GPU）环境需要通过算力卡获得。目前算力卡是通过邀请用户的方式获得。
MegStudio 开发环境是基于 JupyterLab，环境关闭之后额外添加的文件会被销毁。使用体验上和 JupyterLab 基本没有区别，如果官方能默认支持一下代码补全就更好了。
友好的中文文档
对照 官方文档 上手的过程当中，很大的一个体会是 MegEngine 的文档做得很用心。
首先从语言上说，中文书写的文档对于国内的开发者来说无疑是很友好的。
官方文档目前分为基础学习和进阶学习两个部分，MegEngine 团队设计了一个循序渐进的入门指引帮助使用者熟悉 MegEngine。
MegEngine 的文档中穿插了对神经网络知识的简要讲解，比如下面是文档中介绍的 BP 过程。
文档代码中的注释也非常详细，介绍了代码释义，和 API 的细节，阅读起来会比较轻松。
支持基于 Module 的网络搭建
MegEngine 提供两种网络搭建方式：基于 functional （提供常见算子）和基于 Module（提供常见网络层）。
基于 Module 的构建方式和 PyTorch 的风格十分相似，下面分别是 MegEngine 和 PyTorch 的 LeNet 实现代码，可以看出总体的写法没有很大的区别，对 PyTorch 用户来说，熟悉 MegEngine 网络搭建压力很小（官方也提到支持导入 PyTorch Module）。
# MegEngine Implementation
class LeNet(M.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = M.Conv2d(1, 6, 5)
self.relu1 = M.ReLU()
self.pool1 = M.MaxPool2d(2, 2)
self.conv2 = M.Conv2d(6, 16, 5)
self.relu2 = M.ReLU()
self.pool2 = M.MaxPool2d(2, 2)
self.fc1 = M.Linear(16 * 5 * 5, 120)
self.relu3 = M.ReLU()
self.fc2 = M.Linear(120, 84)
self.relu4 = M.ReLU()
self.classifer = M.Linear(84, 10)
def forward(self, x):
x = self.pool1(self.relu1(self.conv1(x)))
x = self.pool2(self.relu2(self.conv2(x)))
x = F.flatten(x, 1)
x = self.relu3(self.fc1(x))
x = self.relu4(self.fc2(x))
x = self.classifer(x)
return x
# PyTorch Implementation
class LeNet5(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 6, 5, padding=2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16*5*5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
x = F.max_pool2d(F.relu(self.conv2(x)), (2, 2))
x = x.view(-1, self.num_flat_features(x))
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
def num_flat_features(self, x):
size = x.size()[1:]
num_features = 1
for s in size:
num_features *= s
return num_features
动态图与静态图的转换
此前的 TensorFlow 框架注重静态图，这对于调试来说比较困难，PyTorch 侧重动态图，对调试开发来说比较便利，但是不利于高效部署。MegEngine 在动态图和静态图的兼容上下了功夫，支持通过用 jit 包中的 trace 装饰器来完成动静转换，这样调试时可以使用动态图，实时输出中间结果，部署时再使用静态图，提升训练和推理速度。这也是 MegEngine 作为工业级框架的重要特性之一。
import megengine.functional as F
from megengine.jit import trace
trace.enabled = True
@trace
def train_func(data, label, *, opt, net):
pred = net(data)
loss = F.cross_entropy_with_softmax(pred, label)
opt.backward(loss)
return pred, loss
train_func(data, label, opt=optimizer, net=le_net)
中文社区
MegEngine 建立了中文语言的 社区，对国内开发者无疑是一大好处，我在站务反馈区提了一个捉虫帖，没想到没几分钟之后管理员就反馈修改完毕，效率可以说是非常高了！
以上就是对 MegEngine 初步的体验，对国内人工智能生态来说，MegEngine 的开源无疑是一件好事，期待未来 MegEngine 和国内深度学习框架的发展！