清华开源Jittor:首个国内高校自研深度学习框架,一键转换PyTorch

作者: 佚名 2020-03-20 14:33:29

 继 Theano、Caffe 之后,又一个由高校主导的深度学习框架开源了,而且还是国产。

深度学习框架越来越多,主导的团队也从高校研究机构渐渐转向了科技巨头。但是,学界在这一领域的力量不容忽视。今日。清华大学开发了一个名为计图(Jittor)的深度学习框架。这一框架有望为深度学习社区提供新的方案,也能够推动深度学习框架国产化的进程。

据悉,计图(Jittor:Just in Time)是一个采用元算子表达神经网络计算单元、完全基于动态编译(Just-in-Time)的深度学习框架,其主要特性为元算子和统一计算图。

据官网介绍说:「研究团队将神经网络所需的基本算子定义为元算子」,元算子非常底层,通过相互融合可以完成复杂的深度学习计算,体现出了易于使用的特点。并且研发团队表示,采用元算子的 Jittor 已超越 Numpy,能够实现更复杂更高效的操作。

而另一层面,统一计算图则是融合了静态计算图和动态计算图的诸多优点,在易于使用的同时,提供高性能的优化。基于元算子开发的深度学习模型,可以被计图实时地自动优化并且运行在指定的硬件上,如 CPU、GPU。

目前 Jittor 已开源,用户可以采用 pip 等方法下载使用。

Jittor 官网:https://cg.cs.tsinghua.edu.cn/jittor/

项目地址:https://github.com/Jittor/jittor

即时、易用、可定制:实现和优化分离的框架

据官网介绍,Jittor 在设计时秉持易用、灵活、即时的设计理念:

  • 易用且可定制:只需要数行代码,就可定义新的算子和模型。
  • 实现与优化分离:可以通过前端接口专注于实现,而实现自动被后端优化。
  • 所有都是即时的:Jittor 的所有代码都是即时编译并且运行的,包括 Jittor 本身。用户可以随时对 Jittor 的所有代码进行修改,并且动态运行。

Jittor 作为全新的深度学习框架,有几大新特性值得关注。首先,Jittor 采用了元算子的概念,将各种基本计算定义为元算子,并通过结合不同元算子,实现深度学习中的各项算子功能。

清华开源Jittor:首个国内高校自研深度学习框架,一键转换PyTorch

元算子融合。

其次,在算子的设置上,团队将元算子的反向传播进行了闭包,即元算子的反向传播也是元算子。这样避免了重复开发。此外,还支持计算任意高阶导数。

清华开源Jittor:首个国内高校自研深度学习框架,一键转换PyTorch

高阶导数及反向传播闭包。

然后,在编程语言上,Jittor 采用了灵活而易用的 Python。用户可以使用它,编写元算子计算的 Python 代码,然后 Jittor

将其动态编译为 C++,实现高性能。

清华开源Jittor:首个国内高校自研深度学习框架,一键转换PyTorch

算子动态编译。

更重要的是,Jittor 内置了优化编译遍(complier pass)功能,和 LLVM 兼容。优化过程会根据运行的硬件而定,以下为已支持的优化编译遍。

清华开源Jittor:首个国内高校自研深度学习框架,一键转换PyTorch

自动优化。

在内存管理上,CPU 和 GPU 内存统一管理。GPU 内存不够的时候会调用 GPU。

清华开源Jittor:首个国内高校自研深度学习框架,一键转换PyTorch

统一内存管理。

而在接口方面,Jittor 同时提供了同步和异步接口,切换不会造成性能损失。从而带来了易用性和高效率。

清华开源Jittor:首个国内高校自研深度学习框架,一键转换PyTorch

高效同步异步接口。

最后值得注意的是,为了帮助 Jittor 的发展,并吸收优秀框架的特性。Jittor 采用了和 PyTorch 相似的模块化接口,并有辅助转换脚本,还有和 PyTorch 一样的 Numpy+pickle 协议,使得两者的模型可以互相加载和使用。

清华开源Jittor:首个国内高校自研深度学习框架,一键转换PyTorch

模型迁移。

有了如此众多的特性,估计用来写模型会有非常好的体验。此外,Jittor 使用 Python 和 C ++编写,也支持 GPU 与 CUDA。

测评结果

Jittor 的性能如何,官方也提供了测评结果进行参考。目前 ResNet、VGG、SSD、DeepLab、LSGAN 等多个网络模型已经在 Jittor 平台实现,可供用户使用。与同类型框架相比,Jittor 在收敛精度一致情况下,推理速度取得了 10%-50% 的性能提升。

清华开源Jittor:首个国内高校自研深度学习框架,一键转换PyTorch

示例代码

Jittor 前端语言为 Python。前端使用模块化的设计,类似于 PyTorch,Keras,后端则使用高性能语言编写,如 CUDA,C++。Jittor 官网目前已提供了示例代码,我们可以来分析一下。

下面的代码演示了如何一步一步使用 Python 代码,从头对一个双层神经网络建模。

  1. import jittor as jt 
  2.  
  3. from jittor import Module 
  4.  
  5. from jittor import nn 
  6.  
  7. class Model(Module): 
  8.  
  9.    def __init__(self): 
  10.  
  11.        self.layer1 = nn.Linear(110
  12.  
  13.        self.relu = nn.Relu() 
  14.  
  15.        self.layer2 = nn.Linear(101
  16.  
  17.    def execute (self,x) : 
  18.  
  19.        x = self.layer1(x) 
  20.  
  21.        x = self.relu(x) 
  22.  
  23.        x = self.layer2(x) 
  24.  
  25.        return x 
  26.  
  27. def get_data(n): # generate random data for training test. 
  28.  
  29.    for i in range(n): 
  30.  
  31.        x = np.random.rand(batch_size, 1
  32.  
  33.        y = x*x 
  34.  
  35.        yield jt.float32(x), jt.float32(y) 
  36.  
  37. model = Model() 
  38.  
  39. learning_rate = 0.1 
  40.  
  41. optim = nn.SGD(model.parameters(), learning_rate) 
  42.  
  43. for i,(x,y) in enumerate(get_data(n)): 
  44.  
  45.    pred_y = model(x) 
  46.  
  47.    loss = ((pred_y - y)**2
  48.  
  49.    loss_mean = loss.mean() 
  50.  
  51.    optim.step (loss_mean) 
  52.  
  53.    print(f"step {i}, loss = {loss_mean.data.sum()}"

以上为实现一个简单的全连接神经网络的代码。可以看到,类似于 PyTorch,只需要导入包,采用类继承的方式定义模型、数据处理(分批等)方式,并设置训练循环即可。从上手难度来看,熟悉 PyTorch 的用户可以直接上手,甚至没有因为变量名不同而带来影响。

研究团队

据官网介绍,Jittor 的开发团队来自清华大学计算机系的图形学实验室,实验室负责人为胡事民教授。

清华开源Jittor:首个国内高校自研深度学习框架,一键转换PyTorch

该实验室成立于 1998 年 3 月,2007 年发展成为清华大学可视媒体研究中心,2010 年获批成为北京市工程技术研究中心,同年和腾讯公司合作成立清华-腾讯互联网创新技术联合实验室,并于 2018 年,成立北京信息科学与技术国家研究中心下的可视媒体智能计算团队。目前有教授 2 名、副教授 3 名、助理研究员 1 名、博士后 4 名和研究生 50 多名。

Jittor 项目的主要研发主力为实验室的梁盾、杨国烨、杨国炜和周文洋等等博士生,此外该项目也得到了清华-腾讯联合实验室的资助和支持。

人工智能 深度学习 技术
上一篇:人工智能机器学习常用算法总结及各个常用算法精确率对比 下一篇:AI如何与医疗行业各个环节深度融合?
评论
取消
暂无评论,快去成为第一个评论的人吧

更多资讯推荐

芯片短缺对人工智能有多大伤害?

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

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

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

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

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

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

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

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

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

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

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

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

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

佚名 ·  1天前
如何使用Auto-Sklearn和Auto-PyTorch实现自动化机器学习

如今,机器学习(ML)正在广泛地影响着商业、工程、以及研究等领域。通常,机器学习水平的进步,与软件和自动化的深入迭代有着密切的关系。

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