【51CTO.com原創稿件】深度學習驅動之下很早創業的中國AI獨角獸曠視，宣布開源自研深度學習框架MegEngine(Brain++核心組件之一)，中文名天元——取自圍棋棋盤中心點的名稱。

[[320459]]

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

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

這樣就可以直接安裝。

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

MNIST數據集中的圖片是28*28的，每張圖被轉化為一個行向量，長度是28*28=784，每一個值代表一個像素點。數據集中共有60000張手寫數據圖片，其中55000張訓練數據，5000張測試數據。

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

標簽是介于0-9之間的數字，用于描述圖片中的數字，轉化為one-hot向量即表示的數字對應的下標為1，其余的值為0。標簽的訓練數據是[55000,10]的數字矩陣。

今天就來拿MNIST來做一個測試。

這是天元進行深度學習進行手寫識別。今天就來測試一下深度學習的GPU場景下訓練速度。

GPU用英偉達顯卡1080ti。對比一下三大框架的訓練速度，代碼實現敏捷度。

下列是天元引擎代碼。

from megengine.data.dataset import MNIST  #導入數據集 
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 
 
 
# 讀取訓練數據并進行預處理 
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),
 # 訓練時一般使用RandomSampler來打亂數據順序 
) 
 
# 實例化網絡 
net = Net() 
 
# SGD優化方法，學習率lr=0.01，動量momentum=0.9 
optimizer = SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4) 
 
total_epochs = 10 # 共運行10個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() # 將參數的梯度置零 
        pred, loss = train_func(batch_data, batch_label, opt=optimizer, net=net) 
        optimizer.step()  # 根據梯度更新參數值 
        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 
 
# 測試數據 
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     此時的x是包含batchsize維度為4的tensor，
即(batchsize，channels，x，y)，x.size(0)指batchsize的值    
把batchsize的值作為網絡的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 相當于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（）函數自帶的索引，從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()   # 所有參數的梯度清零 
        loss.backward()         #即反向傳播求梯度 
        optimizer.step()        #調用optimizer進行梯度下降更新參數 
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() 

測試時間，在1080ti，Ubuntu 18.04的系統下，訓練時間如下。

MegEngine 45.3429

TensorFlow 87.3634

PyTorch 68.8535

弄完代碼測試完了以后，對于天元深度學習框架與TensorFlow與PyTorch進行了一下小節。

1.易用性

易用性上，MegEngine與PyTorch最為簡潔，代碼比TensorFlow要簡潔優化的多，并且無縫兼容PyTorch。TensorFlow較為復雜，學習難度較大。如果要快速上手，用MegEngine也不錯。

TensorFlow是比較不友好的，與Python等語言差距很大，有點像基于一種語言重新定義了一種編程語言，并且在調試的時候比較復雜。每次版本的更新，TensorFlow的各種接口經常會有很大幅度的改變，這也大大增加了對其的學習時間。

PyTorch支持動態計算圖，追求盡量少的封裝，代碼簡潔易讀，應用十分靈活，接口沿用Torch，具有很強的易用性，同時可以很好的利用主語言Python的各種優勢。

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

2.速度

曠世對寫出的天元深度學習框架做了許多優化，在我的顯卡1080ti的狀態MegEngine >PyTorch>TensorFlow。就單機訓練速度而言，MegEngine速度相當快，并且用的內存最小。據說MegEngine采用靜態內存分配，做到了極致的內存優化。這一點上來說是相當難得的。

3.算子數量

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

4.開源模型數據集案例

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

5.開源社區

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

6.靈活性

TensorFlow主要支持靜態計算圖的形式，計算圖的結構比較直觀，但是在調試過程中十分復雜與麻煩，一些錯誤更加難以發現。但是在2017年底發布了動態圖機制Eager Execution，加入對于動態計算圖的支持，但是目前依舊采用原有的靜態計算圖形式為主。TensorFlow擁有TensorBoard應用，可以監控運行過程，可視化計算圖。

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

以表格的形式對比如下： 

注：分數越高越好

總結下MegEngine 有下列優點。

1. 上手簡單
2. 兼容PyTorch
3. 支持靜態圖/動態圖兩種機制
4. 中文教程詳細
5. 中文社區，有利于英文不好的童鞋。

也有一些不足有待改進。

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

希望MegEngine為中國的人工智能產業做出自己的貢獻，打造中國智造的人工智能深度學習引擎，為中華民族崛起貢獻自己的光輝力量，我們也可以好好用用國產深度學習引擎，打造我們偉大的祖國成為AI強國。

作者介紹：

尹成，清華碩士，微軟人工智能領域全球最有價值專家，精通C/C++，Python，Go，TensorFlow,擁有15年編程經驗與5年的教學經驗，資深軟件架構師，Intel軟件技術專家 ，具備深厚的項目管理經驗以及研發經驗，擁有兩項人工智能發明專利。

51CTO的世界五百強算法課

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

【51CTO原創稿件，合作站點轉載請注明原文作者和出處為51CTO.com】