概况

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