分享好友 资讯首页 频道列表

比用Pytorch框架快200倍!0.76秒后,笔记本上的CNN就搞定了MNIST

2021-08-13 00:006050

本文经AI新媒体量子位(公众号ID:QbitAI)授权转载,转载请联系出处。

在MNIST上进行训练,可以说是计算机视觉里的“Hello World”任务了。

而如果使用PyTorch的标准代码训练CNN,一般需要3分钟左右。

但现在,在一台笔记本电脑上就能将时间缩短200多倍。

速度直达0.76秒!

那么,到底是如何仅在一次epoch的训练中就达到99%的准确率的呢?

八步提速200倍

这是一台装有GeForce GTX 1660 Ti GPU的笔记本。

我们需要的还有Python3.x和Pytorch 1.8。

先下载数据集进行训练,每次运行训练14个epoch。

这时两次运行的平均准确率在测试集上为99.185%,平均运行时间为2min 52s ± 38.1ms。

接下来,就是一步一步来减少训练时间:

比用Pytorch框架快200倍!0.76秒后,笔记本上的CNN就搞定了MNIST一、提前停止训练

在经历3到5个epoch,测试准确率达到99%时就提前停止训练。

这时的训练时间就减少了1/3左右,达到了57.4s±6.85s。

二、缩小网络规模,采用正则化的技巧来加快收敛速度

具体的,在第一个conv层之后添加一个2x2的最大采样层(max pool layer),将全连接层的参数减少4倍以上。

然后再将2个dropout层删掉一个。

这样,需要收敛的epoch数就降到了3个以下,训练时间也减少到30.3s±5.28s。

三、优化数据加载

使用data_loader.save_data(),将整个数据集以之前的处理方式保存到磁盘的一个pytorch数组中。

也就是不再一次一次地从磁盘上读取数据,而是将整个数据集一次性加载并保存到GPU内存中。

这时,我们只需要一次epoch,就能将平均训练时间下降到7.31s ± 1.36s。

四、增加Batch Size

将Batch Size从64增加到128,平均训练时间减少到4.66s ± 583ms。

五、提高学习率

使用Superconvergence来代替指数衰减。

在训练开始时学习率为0,到中期线性地最高值(4.0),再慢慢地降到0。

比用Pytorch框架快200倍!0.76秒后,笔记本上的CNN就搞定了MNIST

这使得我们的训练时间下降到3.14s±4.72ms。

六、再次增加Batch Size、缩小缩小网络规模

重复第二步,将Batch Size增加到256。

重复第四步,去掉剩余的dropout层,并通过减少卷积层的宽度来进行补偿。

最终将平均时间降到1.74s±18.3ms。

七、最后的微调

首先,将最大采样层移到线性整流函数(ReLU)激活之前。

然后,将卷积核大小从3增加到5.

最后进行超参数调整:

使学习率为0.01(默认为0.001),beta1为0.7(默认为0.9),bata2为0.9(默认为0.999)。

到这时,我们的训练已经减少到一个epoch,在762ms±24.9ms的时间内达到了99.04%的准确率。

比用Pytorch框架快200倍!0.76秒后,笔记本上的CNN就搞定了MNIST

“这只是一个Hello World案例”

对于这最后的结果,有人觉得司空见惯:

优化数据加载时间,缩小模型尺寸,使用ADAM而不是SGD等等,都是常识性的事情。

我想没有人会真的费心去加速运行MNIST,因为这是机器学习中的“Hello World”,重点只是像你展示最小的关键值,让你熟悉这个框架——事实上3分钟也并不长吧。

比用Pytorch框架快200倍!0.76秒后,笔记本上的CNN就搞定了MNIST

而也有网友觉得,大多数人的工作都不在像是MNIST这样的超级集群上。因此他表示:

我所希望的是工作更多地集中在真正最小化训练时间方面。

比用Pytorch框架快200倍!0.76秒后,笔记本上的CNN就搞定了MNIST

GitHub:
https://github.com/tuomaso/train_mnist_fast

 

反对 0
举报 0
收藏 0
打赏 0
评论 0
新AI框架助力宇宙学:从爆炸恒星中解锁更多宇宙膨胀与暗能量之谜
新AI框架助力宇宙学:从爆炸恒星中解锁更多宇宙膨胀与暗能量之谜

0评论2026-06-111090

新AI框架助力宇宙学:借爆炸恒星图像精准测距,解锁暗能量奥秘
新AI框架助力宇宙学:借爆炸恒星图像精准测距,解锁暗能量奥秘

0评论2026-06-041996

苹果iOS 26.6 Beta 1更新:屏蔽联系人达2万上限提醒,地图添新安全框架
苹果iOS 26.6 Beta 1更新:屏蔽联系人达2万上限提醒,地图添新安全框架

0评论2026-05-272173

国产AI突破:全球首个AI自研预训练框架诞生,大模型进入新纪元
国产AI突破:全球首个AI自研预训练框架诞生,大模型进入新纪元

0评论2026-05-261827

神舟二十三号发射成功;类脑导航框架获突破;智能显微镜助力科研新进展
神舟二十三号发射成功;类脑导航框架获突破;智能显微镜助力科研新进展

0评论2026-05-263017

清华长三角研究院推出RWAI框架及RWAI Arena 助力AI真实场景高效落地
清华长三角研究院推出RWAI框架及RWAI Arena 助力AI真实场景高效落地

0评论2026-05-262957

DeepSeek“识图模式”大范围上线,以创新框架开启图文交互新体验
DeepSeek“识图模式”大范围上线,以创新框架开启图文交互新体验

0评论2026-05-092341

DeepSeek发布多模态大模型技术报告 创新框架突破空间参照瓶颈
DeepSeek发布多模态大模型技术报告 创新框架突破空间参照瓶颈

0评论2026-05-012871

DeepSeek发布多模态大模型技术报告,创新框架突破空间参照瓶颈
DeepSeek发布多模态大模型技术报告,创新框架突破空间参照瓶颈

0评论2026-05-01865

检索增强的三维空间新框架:为机器人操作装上“空间智慧大脑”
检索增强的三维空间新框架:为机器人操作装上“空间智慧大脑”

0评论2026-05-011620