体验中国自主知识产权天元深度学习引擎与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公里内目标识别
评论
取消
暂无评论,快去成为第一个评论的人吧

更多资讯推荐

科技史上超炸组合「元宇宙+脑机接口」,离我们还有多远?

技术推进门槛很硬,资本市场敬而远之,伦理问题纠缠不清。脑机接口距离元宇宙还差了几个"VR元年"?

董子博 ·  1天前
芯片短缺对人工智能有多大伤害?

全球芯片短缺的影响随处可见:从家用电器价格上涨到消费电子设备产能受限。急需处理器的人工智能应用也正面临压力,不过SambaNova CEO表示,硬件本身并不是人工智能成功的决定性因素。

AI时代前沿 ·  1天前
2021年人工智能与自动化的发展趋势

自动化正在成为许多市场的首要任务,特别是随着远程工作的增长和劳动力短缺改变了传统的工作模式,很多企业转而采用更可持续的自动化解决方案。

Shelby Hiter ·  2天前
人工智能能否帮助金融行业有效应对勒索软件?

现在是金融机构安全意识进一步发展的时候了——这意味着要超越试图阻止勒索软件突破防火墙的预防性方法,专注于用能够检测和阻止攻击的工具武装自己。

Garry Veale ·  2天前
AI视频分析技术是如何工作的?原理是什么?

实时 AI 视频分析是一种基于人工智能的技术,可分析视频流以检测特定行为和事件的展开。这种类型的系统通过人工智能机器学习引擎检查来自监控摄像头的视频流来进行相关工作。该引擎使用一系列算法和程序来理解数据,并将数据转换为可理解的、有意义的信息。

EasyNVR ·  2天前
AI能成科学家的工具人?Nature采访五位顶尖学者:学会写代码,降低期望

AI 对于其他领域的科研来说是一个极其好用的工具,DNA测序、天文地理甚至艺术领域都必须要用到AI 模型来提供灵感。最近Nature 采访了五位跨领域的专家,听听他们对于AI 工具人有什么想法?

佚名 ·  2天前
人工智能与云计算正加速形成应用生态

人工智能在赋能生产力升级,推动各行业完成智能化转型和新旧动能转换的进程中发挥着重要作用。同时,人工智能技术也已经广泛应用在金融、教育、医疗、能源、消费、工业等各行业多场景之中。

佚名 ·  2天前
Facebook在ICCV 2021 发布两个3D模型,自监督才是终极答案?

在 ICCV 2021 上,Facebook AI提出了两个新模型3DETR和DepthContrast,这两个互补的新模型可促进3D理解并更容易上手。

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