【51CTO.com原創(chuàng)稿件】深度學(xué)習(xí)驅(qū)動之下很早創(chuàng)業(yè)的中國AI獨(dú)角獸曠視，宣布開源自研深度學(xué)習(xí)框架MegEngine(Brain++核心組件之一)，中文名天元——取自圍棋棋盤中心點(diǎn)的名稱。

[[320459]]

今天就來帶大家體驗(yàn)一下天元深度學(xué)習(xí)框架，安裝天元深度學(xué)習(xí)引擎需要Linux，目前天元只支持Linux，對于熟悉Windows的人員還不夠友好。需要在Windows的sublinux中安裝。除非安裝雙系統(tǒng)，Linux裸機(jī)，否則可能支持不了GPU深度學(xué)習(xí)加速。這一點(diǎn)還不夠足夠方便。但是國產(chǎn)的初創(chuàng)引擎可以理解，后面一定會改進(jìn)的。

pip3 install megengine -f https://megengine.org.cn/whl/mge.html

這樣就可以直接安裝。

手寫識別數(shù)據(jù)集的官網(wǎng)在這里，主要是手寫識別的系列數(shù)據(jù)集。http://yann.lecun.com/exdb/mnist/index.html,我們可以從這里下載數(shù)據(jù)集。

MNIST數(shù)據(jù)集中的圖片是28*28的，每張圖被轉(zhuǎn)化為一個(gè)行向量，長度是28*28=784，每一個(gè)值代表一個(gè)像素點(diǎn)。數(shù)據(jù)集中共有60000張手寫數(shù)據(jù)圖片，其中55000張訓(xùn)練數(shù)據(jù)，5000張測試數(shù)據(jù)。

在MNIST中，mnist.train.images是一個(gè)形狀為[55000, 784]的張量，其中的第一個(gè)維度是用來索引圖片，第二個(gè)維度是圖片中的像素。MNIST數(shù)據(jù)集包含有三部分，訓(xùn)練數(shù)據(jù)集，驗(yàn)證數(shù)據(jù)集，測試數(shù)據(jù)集(mnist.validation)。

標(biāo)簽是介于0-9之間的數(shù)字，用于描述圖片中的數(shù)字，轉(zhuǎn)化為one-hot向量即表示的數(shù)字對應(yīng)的下標(biāo)為1，其余的值為0。標(biāo)簽的訓(xùn)練數(shù)據(jù)是[55000,10]的數(shù)字矩陣。

今天就來拿MNIST來做一個(gè)測試。

這是天元進(jìn)行深度學(xué)習(xí)進(jìn)行手寫識別。今天就來測試一下深度學(xué)習(xí)的GPU場景下訓(xùn)練速度。

GPU用英偉達(dá)顯卡1080ti。對比一下三大框架的訓(xùn)練速度，代碼實(shí)現(xiàn)敏捷度。

下列是天元引擎代碼。

from megengine.data.dataset import MNIST  #導(dǎo)入數(shù)據(jù)集 
train_dataset = MNIST(root="./dataset/MNIST", train=True, download=True) 
test_dataset = MNIST(root="./dataset/MNIST", train=False, download=False) 
import megengine.module as M 
import megengine.functional as F 
 
class Net(M.Module): 
    def __init__(self): 
        super().__init__() 
        self.conv0 = M.Conv2d(1, 20, kernel_size=5, bias=False) 
        self.bn0 = M.BatchNorm2d(20) 
        self.relu0 = M.ReLU() 
        self.pool0 = M.MaxPool2d(2) 
        self.conv1 = M.Conv2d(20, 20, kernel_size=5, bias=False) 
        self.bn1 = M.BatchNorm2d(20) 
        self.relu1 = M.ReLU() 
        self.pool1 = M.MaxPool2d(2) 
        self.fc0 = M.Linear(500, 64, bias=True) 
        self.relu2 = M.ReLU() 
        self.fc1 = M.Linear(64, 10, bias=True) 
 
    def forward(self, x): 
        x = self.conv0(x) 
        x = self.bn0(x) 
        x = self.relu0(x) 
        x = self.pool0(x) 
        x = self.conv1(x) 
        x = self.bn1(x) 
        x = self.relu1(x) 
        x = self.pool1(x) 
        x = F.flatten(x, 1) 
        x = self.fc0(x) 
        x = self.relu2(x) 
        x = self.fc1(x) 
        return x 
from megengine.jit import trace 
 
 
@trace(symbolic=True) 
def train_func(data, label, *, opt, net): 
    net.train() 
    pred = net(data) 
    loss = F.cross_entropy_with_softmax(pred, label) 
    opt.backward(loss) 
    return pred, loss 
 
 
@trace(symbolic=True) 
def eval_func(data, label, *, net): 
    net.eval() 
    pred = net(data) 
    loss = F.cross_entropy_with_softmax(pred, label) 
    return pred, loss 
import time 
import numpy as np 
 
import megengine as mge 
from megengine.optimizer import SGD 
from megengine.data import DataLoader 
from megengine.data.transform import ToMode, Pad, Normalize, Compose 
from megengine.data.sampler import RandomSampler 
 
 
# 讀取訓(xùn)練數(shù)據(jù)并進(jìn)行預(yù)處理 
dataloader = DataLoader( 
    train_dataset, 
    transform=Compose([ 
        Normalize(mean=0.1307*255, std=0.3081*255), 
        Pad(2), 
        ToMode('CHW'), 
    ]), 
    sampler=RandomSampler(dataset=train_dataset, batch_size=64),
 # 訓(xùn)練時(shí)一般使用RandomSampler來打亂數(shù)據(jù)順序 
) 
 
# 實(shí)例化網(wǎng)絡(luò) 
net = Net() 
 
# SGD優(yōu)化方法，學(xué)習(xí)率lr=0.01，動量momentum=0.9 
optimizer = SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4) 
 
total_epochs = 10 # 共運(yùn)行10個(gè)epoch 
for epoch in range(total_epochs): 
    total_loss = 0 
    for step, (batch_data, batch_label) in enumerate(dataloader): 
        batch_label = batch_label.astype(np.int32) 
        optimizer.zero_grad() # 將參數(shù)的梯度置零 
        pred, loss = train_func(batch_data, batch_label, opt=optimizer, net=net) 
        optimizer.step()  # 根據(jù)梯度更新參數(shù)值 
        total_loss += loss.numpy().item() 
    print("epoch: {}, loss {}".format(epoch, total_loss/len(dataloader))) 
mge.save(net.state_dict(), 'mnist_net.mge') 
net = Net() 
state_dict = mge.load('mnist_net.mge') 
net.load_state_dict(state_dict) 
from megengine.data.sampler import SequentialSampler 
 
# 測試數(shù)據(jù) 
test_sampler = SequentialSampler(test_dataset, batch_size=500) 
dataloader_test = DataLoader( 
    test_dataset, 
    sampler=test_sampler, 
    transform=Compose([ 
        Normalize(mean=0.1307*255, std=0.3081*255), 
        Pad(2), 
        ToMode('CHW'), 
    ]), 
) 
 
correct = 0 
total = 0 
for idx, (batch_data, batch_label) in enumerate(dataloader_test): 
    batch_label = batch_label.astype(np.int32) 
    pred, loss = eval_func(batch_data, batch_label, net=net) 
    predicted = F.argmax(pred, axis=1) 
    correct += (predicted == batch_label).sum().numpy().item() 
    total += batch_label.shape[0] 
print("correct: {}, total: {}, accuracy: {}".format(correct, total, float(correct) / total)) 

這是TensorFlow版本。

import tensorflow as tf 
import numpy as np 
from tensorflow.examples.tutorials.mnist import input_data 
import matplotlib.pyplot as plt 
mnist = input_data.read_data_sets('MNIST_data', one_hot=True) 
tf.reset_default_graph() 
x = tf.placeholder(tf.float32, [None, 784]) 
y = tf.placeholder(tf.float32, [None, 10]) 
w = tf.Variable(tf.random_normal([784, 10])) 
b = tf.Variable(tf.zeros([10])) 
pred = tf.matmul(x, w) + b 
pred = tf.nn.softmax(pred) 
cost = tf.reduce_mean(-tf.reduce_sum(y * tf.log(pred), reduction_indices=1)) 
learning_rate = 0.01 
optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(cost) 
training_epochs = 25 
batch_size = 100 
display_step = 1 
save_path = 'model/' 
saver = tf.train.Saver() 
with tf.Session() as sess: 
    sess.run(tf.global_variables_initializer()) 
    for epoch in range(training_epochs): 
        avg_cost = 0 
        total_batch = int(mnist.train.num_examples/batch_size) 
        for i in range(total_batch): 
            batch_xs, batch_ys = mnist.train.next_batch(batch_size) 
            _, c = sess.run([optimizer, cost], feed_dict={x:batch_xs, y:batch_ys}) 
            avg_cost += c / total_batch 
 
        if (epoch + 1) % display_step == 0: 
            print('epoch= ', epoch+1, ' cost= ', avg_cost) 
    print('finished') 
    correct_prediction = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1)) 
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32)) 
    print('accuracy: ', accuracy.eval({x:mnist.test.images, y:mnist.test.labels})) 
    save = saver.save(sess, save_path=save_path+'mnist.cpkt') 
print(" starting 2nd session ...... ") 
with tf.Session() as sess: 
    sess.run(tf.global_variables_initializer()) 
    saver.restore(sess, save_path=save_path+'mnist.cpkt') 
    correct_prediction = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1)) 
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32)) 
    print('accuracy: ', accuracy.eval({x: mnist.test.images, y: mnist.test.labels})) 
    output = tf.argmax(pred, 1) 
    batch_xs, batch_ys = mnist.test.next_batch(2) 
    outputval= sess.run([output], feed_dict={x:batch_xs, y:batch_ys}) 
    print(outputval) 
    im = batch_xs[0] 
    im = im.reshape(-1, 28) 
    plt.imshow(im, cmap='gray') 
    plt.show() 
    im = batch_xs[1] 
    im = im.reshape(-1, 28) 
    plt.imshow(im, cmap='gray') 
    plt.show() 

這是PyTorch版本。

import torch 
import torch.nn as nn 
import torch.nn.functional as F 
import torch.optim as optim 
from torchvision import datasets, transforms 
from torch.autograd import Variable 
batch_size = 64 
train_dataset = datasets.MNIST(root='./data/', 
                               train=True, 
                               transform=transforms.ToTensor(), 
                               download=True) 
 
test_dataset = datasets.MNIST(root='./data/', 
                              train=False, 
                              transform=transforms.ToTensor()) 
 
# Data Loader (Input Pipeline) 
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, 
                                           batch_size=batch_size, 
                                           shuffle=True) 
 
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, 
                                          batch_size=batch_size, 
                                          shuffle=False) 
class Net(nn.Module): 
    def __init__(self): 
        super(Net, self).__init__() 
        # 輸入1通道，輸出10通道，kernel 5*5 
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=10, kernel_size=5) 
        self.conv2 = nn.Conv2d(10, 20, 5) 
        self.conv3 = nn.Conv2d(20, 40, 3) 
 
        self.mp = nn.MaxPool2d(2) 
        # fully connect 
        self.fc = nn.Linear(40, 10)#（in_features, out_features） 
    def forward(self, x): 
        # in_size = 64 
        in_size = x.size(0) # one batch     此時(shí)的x是包含batchsize維度為4的tensor，
即(batchsize，channels，x，y)，x.size(0)指batchsize的值    
把batchsize的值作為網(wǎng)絡(luò)的in_size 
        # x: 64*1*28*28 
        x = F.relu(self.mp(self.conv1(x))) 
        # x: 64*10*12*12  feature map =[(28-4)/2]^2=12*12 
        x = F.relu(self.mp(self.conv2(x))) 
        # x: 64*20*4*4 
        x = F.relu(self.mp(self.conv3(x))) 
 
        x = x.view(in_size, -1) # flatten the tensor 相當(dāng)于resharp 
        # print(x.size()) 
        # x: 64*320 
        x = self.fc(x) 
        # x:64*10 
        # print(x.size()) 
        return F.log_softmax(x)  #64*10 
model = Net() 
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5) 
def train(epoch): 
    for batch_idx, (data, target) in enumerate(train_loader):
#batch_idx是enumerate（）函數(shù)自帶的索引，從0開始 
        # data.size():[64, 1, 28, 28] 
        # target.size():[64] 
        output = model(data) 
        #output:64*10 
        loss = F.nll_loss(output, target) 
        if batch_idx % 200 == 0: 
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( 
                epoch, batch_idx * len(data), len(train_loader.dataset), 
                100. * batch_idx / len(train_loader), loss.data[0])) 
        optimizer.zero_grad()   # 所有參數(shù)的梯度清零 
        loss.backward()         #即反向傳播求梯度 
        optimizer.step()        #調(diào)用optimizer進(jìn)行梯度下降更新參數(shù) 
def test(): 
    test_loss = 0 
    correct = 0 
    for data, target in test_loader: 
        data, target = Variable(data, volatile=True), Variable(target) 
        output = model(data) 
        # sum up batch loss 
        test_loss += F.nll_loss(output, target, size_average=False).data[0] 
        # get the index of the max log-probability 
        pred = output.data.max(1, keepdim=True)[1] 
        print(pred) 
        correct += pred.eq(target.data.view_as(pred)).cpu().sum() 
    test_loss /= len(test_loader.dataset) 
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( 
        test_loss, correct, len(test_loader.dataset), 
        100. * correct / len(test_loader.dataset))) 
for epoch in range(1, 10): 
    train(epoch) 
    test() 

測試時(shí)間，在1080ti，Ubuntu 18.04的系統(tǒng)下，訓(xùn)練時(shí)間如下。

MegEngine 45.3429

TensorFlow 87.3634

PyTorch 68.8535

弄完代碼測試完了以后，對于天元深度學(xué)習(xí)框架與TensorFlow與PyTorch進(jìn)行了一下小節(jié)。

1.易用性

易用性上，MegEngine與PyTorch最為簡潔，代碼比TensorFlow要簡潔優(yōu)化的多，并且無縫兼容PyTorch。TensorFlow較為復(fù)雜，學(xué)習(xí)難度較大。如果要快速上手，用MegEngine也不錯(cuò)。

TensorFlow是比較不友好的，與Python等語言差距很大，有點(diǎn)像基于一種語言重新定義了一種編程語言，并且在調(diào)試的時(shí)候比較復(fù)雜。每次版本的更新，TensorFlow的各種接口經(jīng)常會有很大幅度的改變，這也大大增加了對其的學(xué)習(xí)時(shí)間。

PyTorch支持動態(tài)計(jì)算圖，追求盡量少的封裝，代碼簡潔易讀，應(yīng)用十分靈活，接口沿用Torch，具有很強(qiáng)的易用性，同時(shí)可以很好的利用主語言Python的各種優(yōu)勢。

對于文檔的詳細(xì)程度，TensorFlow具備十分詳盡的官方文檔，查找起來十分方便，同時(shí)保持很快的更新速度，但是條理不是很清晰，教程眾多。MegEngine與PyTorch相對而言，條理清晰。PyTorch案例豐富度略多，MegEngine還在建設(shè)，相信將來一定會很多。整體而言，易用性上MegEngine與PyTorch差不多。TensorFlow的易用性最差。 MegEngine可以無縫借用PyTorch的所有案例，這一點(diǎn)兼容性上做的非常贊。

2.速度

曠世對寫出的天元深度學(xué)習(xí)框架做了許多優(yōu)化，在我的顯卡1080ti的狀態(tài)MegEngine >PyTorch>TensorFlow。就單機(jī)訓(xùn)練速度而言，MegEngine速度相當(dāng)快，并且用的內(nèi)存最小。據(jù)說MegEngine采用靜態(tài)內(nèi)存分配，做到了極致的內(nèi)存優(yōu)化。這一點(diǎn)上來說是相當(dāng)難得的。

3.算子數(shù)量

這一點(diǎn)TensorFlow無疑是最大的贏家，提供的Python API達(dá)到8000多個(gè)(參見https://tensorflow.google.cn/api_docs/python)，基本上不存在用戶找不到的算子，所有的算法都可以用TensorFlow算子拼出來。不過API過多也是個(gè)負(fù)擔(dān)，又是low level又是high level，容易把用戶整暈。PyTorch的算子其次，優(yōu)化較TensorFlow較多。天元深度學(xué)習(xí)框架算子數(shù)量目前最少，但是優(yōu)化較好，算子質(zhì)量較高。其中任意圖在訓(xùn)練MNIST上就有很大幫助。目前TensorFlow與PyTorch都不支持。

4.開源模型數(shù)據(jù)集案例

目前TensorFlow>PyTorch>MegEngine，這一塊TensorFlow時(shí)間最久，案例最多，PyTorch其次，MegEngine最少。

5.開源社區(qū)

目前MegEngine的開源社區(qū)最小，不過也在籌備組建中，盡早加入也能成為類似PyTorch,TensorFlow社區(qū)的大牛。如果招聘天元深度學(xué)習(xí)框架的人越多，可能促進(jìn)學(xué)習(xí)的人越多。目前國產(chǎn)還在起步路上，但是非常看好天元的深度學(xué)習(xí)框架的未來。

6.靈活性

TensorFlow主要支持靜態(tài)計(jì)算圖的形式，計(jì)算圖的結(jié)構(gòu)比較直觀，但是在調(diào)試過程中十分復(fù)雜與麻煩，一些錯(cuò)誤更加難以發(fā)現(xiàn)。但是在2017年底發(fā)布了動態(tài)圖機(jī)制Eager Execution，加入對于動態(tài)計(jì)算圖的支持，但是目前依舊采用原有的靜態(tài)計(jì)算圖形式為主。TensorFlow擁有TensorBoard應(yīng)用，可以監(jiān)控運(yùn)行過程，可視化計(jì)算圖。

PyTorch為動態(tài)計(jì)算圖的典型代表，便于調(diào)試，并且高度模塊化，搭建模型十分方便，同時(shí)具備極其優(yōu)秀的GPU支持，數(shù)據(jù)參數(shù)在CPU與GPU之間遷移十分靈活。MegEngine無縫兼容PyTorch，并且具備動態(tài)圖這樣的優(yōu)點(diǎn)，在分布式部署，模型推斷上有較大優(yōu)勢。MegEngine的靈活性相比TensorFlow有較大優(yōu)勢。

以表格的形式對比如下： 

注：分?jǐn)?shù)越高越好

總結(jié)下MegEngine 有下列優(yōu)點(diǎn)。

1. 上手簡單
2. 兼容PyTorch
3. 支持靜態(tài)圖/動態(tài)圖兩種機(jī)制
4. 中文教程詳細(xì)
5. 中文社區(qū)，有利于英文不好的童鞋。

也有一些不足有待改進(jìn)。

1. 支持 Windows/Mac及其他 Linux發(fā)行版。目前只支持 Ubuntu 16.04 及以上，這個(gè)確實(shí)有點(diǎn)少。而且應(yīng)該還不支持 Docker等安裝方式。
2. 更多的案例庫與模型庫。
3. 盡量早日進(jìn)駐各大高校，有教材與教程，建立中國自己的AI人才梯隊(duì)。

希望MegEngine為中國的人工智能產(chǎn)業(yè)做出自己的貢獻(xiàn)，打造中國智造的人工智能深度學(xué)習(xí)引擎，為中華民族崛起貢獻(xiàn)自己的光輝力量，我們也可以好好用用國產(chǎn)深度學(xué)習(xí)引擎，打造我們偉大的祖國成為AI強(qiáng)國。

作者介紹：

尹成，清華碩士，微軟人工智能領(lǐng)域全球最有價(jià)值專家，精通C/C++，Python，Go，TensorFlow,擁有15年編程經(jīng)驗(yàn)與5年的教學(xué)經(jīng)驗(yàn)，資深軟件架構(gòu)師，Intel軟件技術(shù)專家 ，具備深厚的項(xiàng)目管理經(jīng)驗(yàn)以及研發(fā)經(jīng)驗(yàn)，擁有兩項(xiàng)人工智能發(fā)明專利。

51CTO的世界五百強(qiáng)算法課

https://edu.51cto.com/sd/0a881

【51CTO原創(chuàng)稿件，合作站點(diǎn)轉(zhuǎn)載請注明原文作者和出處為51CTO.com】