概况
天元(英文名: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 和国内深度学习框架的发展!
发表回复