当前位置: 首页 > 科技观察

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

时间:2023-03-22 16:18:00 科技观察

继Theano、Caffe之后,又一高校主导的深度学习框架开源,目前仍为国产.深度学习框架越来越多,领导团队也逐渐从高校研究机构转向科技巨头。但是,该领域学术界的力量也不容忽视。今天。清华大学开发了一个名为Jittor的深度学习框架。该框架有望为深度学习社区提供新的解决方案,也能推动深度学习框架国产化进程。据悉,极兔(Jittor:JustinTime)是一个使用元算子表达神经网络计算单元,完全基于动态编译(Just-in-Time)的深度学习框架。它的主要特点是元算子和统一的计算图。.据官网介绍,“研究团队将神经网络所需的基本算子定义为元算子”。元算子非常底层,通过相互融合可以完成复杂的深度学习计算,体现了易用性的特点。并且研发团队表示,使用元算子的Jittor已经超越了Numpy,可以实现更复杂、更高效的运算。另一方面,统一计算图结合了静态计算图和动态计算图的诸多优点,在提供高性能优化的同时还易于使用。基于元算子开发的深度学习模型可以实时自动优化并运行在指定的硬件上,如CPU和GPU。Jittor目前是开源的,用户可以通过pip等方式下载使用。Jittor官网:https://cg.cs.tsinghua.edu.cn/jittor/项目地址:https://github.com/Jittor/jittor即时、易用、可定制:实现与优化分离的框架据官网介绍,Jittor在设计时秉承易用、灵活、实时的设计理念:易用可定制:只需几行代码即可定义新的算子和模型。实现与优化分离:可以通过前端接口专注实现,后端自动优化实现。Alljust-in-time:Jittor中的所有代码都是即时编译和运行的,包括Jittor本身。用户可以随时修改Jittor的所有代码并动态运行。Jittor作为一个全新的深度学习框架,有几个值得关注的新特性。首先,Jittor采用元算子的概念,将各种基础计算定义为元算子,通过组合不同的元算子实现深度学习中的各种算子功能。元算子融合。其次,在算子的设置上,团队关闭了元算子的反向传播,即元算子的反向传播也是元算子。这避免了重复开发。此外,还支持计算任意高阶导数。高阶导数和反向传播闭包。然后,在编程语言方面,Jittor采用了灵活易用的Python。用户可以用它编写元算子计算的Python代码,然后Jittor会动态编译成C++实现高性能。算子动态编译。更重要的是,Jittor内置了complierpass功能,兼容LLVM。优化过程将取决于正在运行的硬件。以下是支持的优化编译过程。自动优化。在内存管理方面,CPU和GPU内存统一管理。当GPU显存不够时,会调用GPU。统一内存管理。在接口方面,Jittor同时提供了同步和异步接口,切换不会造成性能损失。结果是易于使用和高效率。高效的同步异步接口。最后但同样重要的是,为了帮助Jittor开发和吸收优秀框架的特性。Jittor采用了类似PyTorch的模块化接口,并有辅助的转换脚本,以及与PyTorch相同的Numpy+pickle协议,使得两者的模型可以相互加载使用。模型迁移。这么多功能,估计写模型会有很好的体验。此外,Jittor是用Python和C++编写的,还支持GPU和CUDA。评测结果展示了Jittor的性能如何,官方的评测结果也提供给大家参考。目前,ResNet、VGG、SSD、DeepLab、LSGAN等多种网络模型已经在Jittor平台上实现并可供用户使用。与同类框架相比,Jittor在收敛精度一致的情况下,在推理速度上实现了10%-50%的性能提升。示例代码Jittor前端语言为Python。前端采用模块化设计,类似于PyTorch和Keras,后端采用高性能语言编写,如CUDA和C++。Jittor官网目前提供了示例代码,我们可以分析一下。以下代码演示了如何使用Python代码从头开始逐步建模两层神经网络。importjittorasjtfromjittoimportModulefromjittorimportnnclassModel(模块):def__init__(self):self.layer1=nn.Linear(1,10)self.relu=nn.Relu()self.layer2=nn.Linear(10,1)defexecute(self,x):x=self.layer1(x)x=self.relu(x)x=self.layer2(x)返回xdefget_data(n):#generaterandomdatafortrainingtest.foriinrange(n):x=np.random.rand(batch_size,1)y=x*xyieldjt.float32(x),jt.float32(y)model=Model()learning_rate=0.1optim=nn.SGD(model.parameters(),learning_rate)fori,(x,y)inenumerate(get_data(n)):pred_y=model(x)loss=((pred_y-y)**2)loss_mean=loss.mean()optim.step(loss_mean)print(f"step{i},loss={loss_mean.data.sum()}")上面是实现一个简单的全连接神经网络的代码。可以看出,与PyTorch类似,只需要通过类继承导入包,定义模型,数据处理(批处理等)方法,设置训练周期即可。从上手难度来看,熟悉PyTorch的用户可以直接上手,甚至不受变量名不同的影响。研究团队官网介绍,Jittor的开发团队来自清华大学计算机系图形学实验室,实验室主任为胡世民教授。实验室成立于1998年3月,2007年发展为清华大学视觉媒体研究中心,2010年获批北京市工程技术研究中心,同年与腾讯合作成立清华-腾讯互联网创新技术联合实验室。并于2018年成立北京国家信息科学技术研究中心视觉媒体智能计算团队。现有教授2人,副教授3人,助理研究员1人,博士后4人,研究生50余人。Jittor项目的主要研发力量是梁盾、杨国业、杨国伟、周文扬等实验室博士生。此外,该项目还得到了清华-腾讯联合实验室的资助和支持。