体验中国自主知识产权天元深度学习引擎与TensorFlow,PyTorch的对比

作者: 尹成 2020-03-31 10:15:42

【51CTO.com原创稿件】深度学习驱动之下很早创业的中国AI独角兽旷视,宣布开源自研深度学习框架MegEngine(Brain++核心组件之一),中文名天元——取自围棋棋盘中心点的名称。

今天就来带大家体验一下天元深度学习框架,安装天元深度学习引擎需要Linux,目前天元只支持Linux,对于熟悉Windows的人员还不够友好。需要在Windows的sublinux中安装。除非安装双系统,Linux裸机,否则可能支持不了GPU深度学习加速。这一点还不够足够方便。但是国产的初创引擎可以理解,后面一定会改进的。

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

这样就可以直接安装。

手写识别数据集的官网在这里,主要是手写识别的系列数据集。https://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。对比一下三大框架的训练速度,代码实现敏捷度。

下列是天元引擎代码。

  1. from megengine.data.dataset import MNIST  #导入数据集 
  2. train_dataset = MNIST(root="./dataset/MNIST", train=True, download=True
  3. test_dataset = MNIST(root="./dataset/MNIST", train=False, download=False
  4. import megengine.module as M 
  5. import megengine.functional as F 
  6.  
  7. class Net(M.Module): 
  8.     def __init__(self): 
  9.         super().__init__() 
  10.         self.conv0 = M.Conv2d(1, 20, kernel_size=5, bias=False
  11.         self.bn0 = M.BatchNorm2d(20) 
  12.         self.relu0 = M.ReLU() 
  13.         self.pool0 = M.MaxPool2d(2) 
  14.         self.conv1 = M.Conv2d(20, 20, kernel_size=5, bias=False
  15.         self.bn1 = M.BatchNorm2d(20) 
  16.         self.relu1 = M.ReLU() 
  17.         self.pool1 = M.MaxPool2d(2) 
  18.         self.fc0 = M.Linear(500, 64, bias=True
  19.         self.relu2 = M.ReLU() 
  20.         self.fc1 = M.Linear(64, 10, bias=True
  21.  
  22.     def forward(self, x): 
  23.         x = self.conv0(x) 
  24.         x = self.bn0(x) 
  25.         x = self.relu0(x) 
  26.         x = self.pool0(x) 
  27.         x = self.conv1(x) 
  28.         x = self.bn1(x) 
  29.         x = self.relu1(x) 
  30.         x = self.pool1(x) 
  31.         x = F.flatten(x, 1) 
  32.         x = self.fc0(x) 
  33.         x = self.relu2(x) 
  34.         x = self.fc1(x) 
  35.         return x 
  36. from megengine.jit import trace 
  37.  
  38.  
  39. @trace(symbolic=True
  40. def train_func(data, label, *, opt, net): 
  41.     net.train() 
  42.     pred = net(data) 
  43.     loss = F.cross_entropy_with_softmax(pred, label) 
  44.     opt.backward(loss) 
  45.     return pred, loss 
  46.  
  47.  
  48. @trace(symbolic=True
  49. def eval_func(data, label, *, net): 
  50.     net.eval() 
  51.     pred = net(data) 
  52.     loss = F.cross_entropy_with_softmax(pred, label) 
  53.     return pred, loss 
  54. import time 
  55. import numpy as np 
  56.  
  57. import megengine as mge 
  58. from megengine.optimizer import SGD 
  59. from megengine.data import DataLoader 
  60. from megengine.data.transform import ToMode, Pad, Normalize, Compose 
  61. from megengine.data.sampler import RandomSampler 
  62.  
  63.  
  64. # 读取训练数据并进行预处理 
  65. dataloader = DataLoader( 
  66.     train_dataset, 
  67.     transform=Compose([ 
  68.         Normalize(mean=0.1307*255, std=0.3081*255), 
  69.         Pad(2), 
  70.         ToMode('CHW'), 
  71.     ]), 
  72.     sampler=RandomSampler(dataset=train_dataset, batch_size=64),
  73.  # 训练时一般使用RandomSampler来打乱数据顺序 
  74.  
  75. # 实例化网络 
  76. net = Net() 
  77.  
  78. # SGD优化方法,学习率lr=0.01,动量momentum=0.9 
  79. optimizer = SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4) 
  80.  
  81. total_epochs = 10 # 共运行10个epoch 
  82. for epoch in range(total_epochs): 
  83.     total_loss = 0 
  84.     for step, (batch_data, batch_label) in enumerate(dataloader): 
  85.         batch_label = batch_label.astype(np.int32) 
  86.         optimizer.zero_grad() # 将参数的梯度置零 
  87.         pred, loss = train_func(batch_data, batch_label, opt=optimizer, net=net) 
  88.         optimizer.step()  # 根据梯度更新参数值 
  89.         total_loss += loss.numpy().item() 
  90.     print("epoch: {}, loss {}".format(epoch, total_loss/len(dataloader))) 
  91. mge.save(net.state_dict(), 'mnist_net.mge'
  92. net = Net() 
  93. state_dict = mge.load('mnist_net.mge'
  94. net.load_state_dict(state_dict) 
  95. from megengine.data.sampler import SequentialSampler 
  96.  
  97. # 测试数据 
  98. test_sampler = SequentialSampler(test_dataset, batch_size=500) 
  99. dataloader_test = DataLoader( 
  100.     test_dataset, 
  101.     sampler=test_sampler, 
  102.     transform=Compose([ 
  103.         Normalize(mean=0.1307*255, std=0.3081*255), 
  104.         Pad(2), 
  105.         ToMode('CHW'), 
  106.     ]), 
  107.  
  108. correct = 0 
  109. total = 0 
  110. for idx, (batch_data, batch_label) in enumerate(dataloader_test): 
  111.     batch_label = batch_label.astype(np.int32) 
  112.     pred, loss = eval_func(batch_data, batch_label, net=net) 
  113.     predicted = F.argmax(pred, axis=1) 
  114.     correct += (predicted == batch_label).sum().numpy().item() 
  115.     total += batch_label.shape[0] 
  116. print("correct: {}, total: {}, accuracy: {}".format(correct, total, float(correct) / total)) 

这是TensorFlow版本。

  1. import tensorflow as tf 
  2. import numpy as np 
  3. from tensorflow.examples.tutorials.mnist import input_data 
  4. import matplotlib.pyplot as plt 
  5. mnist = input_data.read_data_sets('MNIST_data', one_hot=True
  6. tf.reset_default_graph() 
  7. x = tf.placeholder(tf.float32, [None, 784]) 
  8. y = tf.placeholder(tf.float32, [None, 10]) 
  9. w = tf.Variable(tf.random_normal([784, 10])) 
  10. b = tf.Variable(tf.zeros([10])) 
  11. pred = tf.matmul(x, w) + b 
  12. pred = tf.nn.softmax(pred) 
  13. cost = tf.reduce_mean(-tf.reduce_sum(y * tf.log(pred), reduction_indices=1)) 
  14. learning_rate = 0.01 
  15. optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(cost) 
  16. training_epochs = 25 
  17. batch_size = 100 
  18. display_step = 1 
  19. save_path = 'model/' 
  20. saver = tf.train.Saver() 
  21. with tf.Session() as sess: 
  22.     sess.run(tf.global_variables_initializer()) 
  23.     for epoch in range(training_epochs): 
  24.         avg_cost = 0 
  25.         total_batch = int(mnist.train.num_examples/batch_size) 
  26.         for i in range(total_batch): 
  27.             batch_xs, batch_ys = mnist.train.next_batch(batch_size) 
  28.             _, c = sess.run([optimizer, cost], feed_dict={x:batch_xs, y:batch_ys}) 
  29.             avg_cost += c / total_batch 
  30.  
  31.         if (epoch + 1) % display_step == 0: 
  32.             print('epoch= ', epoch+1, ' cost= ', avg_cost) 
  33.     print('finished'
  34.     correct_prediction = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1)) 
  35.     accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32)) 
  36.     print('accuracy: ', accuracy.eval({x:mnist.test.images, y:mnist.test.labels})) 
  37.     save = saver.save(sess, save_path=save_path+'mnist.cpkt'
  38. print(" starting 2nd session ...... "
  39. with tf.Session() as sess: 
  40.     sess.run(tf.global_variables_initializer()) 
  41.     saver.restore(sess, save_path=save_path+'mnist.cpkt'
  42.     correct_prediction = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1)) 
  43.     accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32)) 
  44.     print('accuracy: ', accuracy.eval({x: mnist.test.images, y: mnist.test.labels})) 
  45.     output = tf.argmax(pred, 1) 
  46.     batch_xs, batch_ys = mnist.test.next_batch(2) 
  47.     outputval= sess.run([output], feed_dict={x:batch_xs, y:batch_ys}) 
  48.     print(outputval) 
  49.     im = batch_xs[0] 
  50.     im = im.reshape(-1, 28) 
  51.     plt.imshow(im, cmap='gray'
  52.     plt.show() 
  53.     im = batch_xs[1] 
  54.     im = im.reshape(-1, 28) 
  55.     plt.imshow(im, cmap='gray'
  56.     plt.show() 

这是PyTorch版本。

  1. import torch 
  2. import torch.nn as nn 
  3. import torch.nn.functional as F 
  4. import torch.optim as optim 
  5. from torchvision import datasets, transforms 
  6. from torch.autograd import Variable 
  7. batch_size = 64 
  8. train_dataset = datasets.MNIST(root='./data/'
  9.                                train=True
  10.                                transform=transforms.ToTensor(), 
  11.                                download=True
  12.  
  13. test_dataset = datasets.MNIST(root='./data/'
  14.                               train=False
  15.                               transform=transforms.ToTensor()) 
  16.  
  17. # Data Loader (Input Pipeline) 
  18. train_loader = torch.utils.data.DataLoader(dataset=train_dataset, 
  19.                                            batch_size=batch_size, 
  20.                                            shuffle=True
  21.  
  22. test_loader = torch.utils.data.DataLoader(dataset=test_dataset, 
  23.                                           batch_size=batch_size, 
  24.                                           shuffle=False
  25. class Net(nn.Module): 
  26.     def __init__(self): 
  27.         super(Net, self).__init__() 
  28.         # 输入1通道,输出10通道,kernel 5*5 
  29.         self.conv1 = nn.Conv2d(in_channels=1, out_channels=10, kernel_size=5) 
  30.         self.conv2 = nn.Conv2d(10, 20, 5) 
  31.         self.conv3 = nn.Conv2d(20, 40, 3) 
  32.  
  33.         self.mp = nn.MaxPool2d(2) 
  34.         # fully connect 
  35.         self.fc = nn.Linear(40, 10)#(in_features, out_features) 
  36.     def forward(self, x): 
  37.         # in_size = 64 
  38.         in_size = x.size(0) # one batch     此时的x是包含batchsize维度为4的tensor,
  39. 即(batchsize,channels,x,y),x.size(0)指batchsize的值    
  40. 把batchsize的值作为网络的in_size 
  41.         # x: 64*1*28*28 
  42.         x = F.relu(self.mp(self.conv1(x))) 
  43.         # x: 64*10*12*12  feature map =[(28-4)/2]^2=12*12 
  44.         x = F.relu(self.mp(self.conv2(x))) 
  45.         # x: 64*20*4*4 
  46.         x = F.relu(self.mp(self.conv3(x))) 
  47.  
  48.         x = x.view(in_size, -1) # flatten the tensor 相当于resharp 
  49.         # print(x.size()) 
  50.         # x: 64*320 
  51.         x = self.fc(x) 
  52.         # x:64*10 
  53.         # print(x.size()) 
  54.         return F.log_softmax(x)  #64*10 
  55. model = Net() 
  56. optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5) 
  57. def train(epoch): 
  58.     for batch_idx, (data, target) in enumerate(train_loader):
  59. #batch_idx是enumerate()函数自带的索引,从0开始 
  60.         # data.size():[64, 1, 28, 28] 
  61.         # target.size():[64] 
  62.         output = model(data) 
  63.         #output:64*10 
  64.         loss = F.nll_loss(output, target) 
  65.         if batch_idx % 200 == 0: 
  66.             print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( 
  67.                 epoch, batch_idx * len(data), len(train_loader.dataset), 
  68.                 100. * batch_idx / len(train_loader), loss.data[0])) 
  69.         optimizer.zero_grad()   # 所有参数的梯度清零 
  70.         loss.backward()         #即反向传播求梯度 
  71.         optimizer.step()        #调用optimizer进行梯度下降更新参数 
  72. def test(): 
  73.     test_loss = 0 
  74.     correct = 0 
  75.     for data, target in test_loader: 
  76.         data, target = Variable(data, volatile=True), Variable(target) 
  77.         output = model(data) 
  78.         # sum up batch loss 
  79.         test_loss += F.nll_loss(output, target, size_average=False).data[0] 
  80.         # get the index of the max log-probability 
  81.         pred = output.data.max(1, keepdim=True)[1] 
  82.         print(pred) 
  83.         correct += pred.eq(target.data.view_as(pred)).cpu().sum() 
  84.     test_loss /= len(test_loader.dataset) 
  85.     print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( 
  86.         test_loss, correct, len(test_loader.dataset), 
  87.         100. * correct / len(test_loader.dataset))) 
  88. for epoch in range(1, 10): 
  89.     train(epoch) 
  90.     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】

人工智能 深度学习 天元
上一篇:跨物种造物,AI将长颈鹿图片转换为鸟,还骗过了人类与机器 下一篇:高科技:美国开发远程人脸识别系统 实现1公里内目标识别
评论
取消
暂无评论,快去成为第一个评论的人吧

更多资讯推荐

从演电影到开车,细数人工智能的五大落地方向

越来越多的企业希望通过利用人工智能的功能来提高其ROI。这篇文章就将带你了解如今人工智能的五大落地方向。

读芯术 ·  1天前
新AI让教师能够快速开发智能辅导系统

通过使用一种采用人工智能的新方法,教师可以通过演示解决某个主题中的问题的几种方法来教计算机……

佚名 ·  1天前
拒绝简单代工,有人工智能的制造业前景更加光明

改革吹风吹大地,各界人士鼓干劲。在工业4.0大背景下,传统产业转型升级的浪潮如约而至,冲击着原有的产业结构与体系。人工智能与工业的深度融合,在无形中塑造着制造业的格局与面貌。

今夕何夕 ·  1天前
AI又一次打败了人类:仅凭一张自拍照片辨别个人性格

一项新研究表明,人工智能(AI)可在“仅凭一张自拍照片辨别个人性格”这件事上再次打败人类。与人类相比,机器能够更好地识别一种特征,通过分析人物面孔来确定物种不同的性格特质。

佚名 ·  1天前
终结重复工作!教你30分钟创建自己的深度学习机器

建立一个深度学习环境是一件很重要的事情。本文讲述使用深度学习 CommunityAMI、TMUX和 Tunneling在EC2为Jupyter Notebooks创建一个新的深度学习服务器。

读芯术 ·  1天前
化身阿凡达,国外小哥开源 AI 实时变脸工具 Avatarify

疫情期间,很多人都在用 Zoom 和 Skype 等软件进行视频会议。国外一位开发小哥觉得这样开会太无聊,于是基于 first-order-model 开发了一个能够把别人的照片套在自己脸上的“变脸”软件 Avatarify,并将其开源。

OSC神秘老司机 ·  1天前
“人脸识别”已衍生出“性格识别”,科技向善还要多久?

近日,《科学报告》期刊上刊登了一篇有关于人脸识别的新技术,俄罗斯研究团队开发了一款新 AI,可“仅凭一张自拍照片辨别个人性格”。

缙霄 ·  1天前
下一站AI:实时服务

随着实时解决方案的增长与人工智能技术的发展,工作负载的日益提升以及非结构化数据的爆炸式增长,数据中心的发展方向正朝着加速计算、存储与网络适应性前进。

佚名 ·  2天前
Copyright©2005-2020 51CTO.COM 版权所有 未经许可 请勿转载