干货|TensorFlow数据量少的时候却占GPU显存比较多

川长思鸟来 2021-11-23 14:46 461阅读 0赞

转自：[https://cloud.tencent.com/developer/article/1365920][https_cloud.tencent.com_developer_article_1365920]

问题背景：

最近在做一个文本多分类的项目，来源于实际的需求场景。具体的情况不多说，但是有一点需要说明的是，场景有多个，每个场景下都有自己的数据，这些数据都是短文本数据。不同的是每个场景中含有的数据量不同。一开始我们做的时候是从数据量最大的场景入手，有107万条训练数据，单词有7万多个，分类效果还不错，不做任何数据预处理，测试集上准确率有94%，这个时候显示的[GPU][]显存是700MB。接着做数据量小一点儿的场景，有70几万条数据，单词有6万多个，发现这个时候的GPU显存有3000多MB。训练时候的参数一模一样。按道理应该单词数多的那个显存比较大才对。而且按照我们的参数计算GPU显存就应该只有几百MB才是正常的。虽然从准确率上看程序应该没问题，但这个问题不解决会让我怀疑自己。

解决思路

先去搜索了一圈，发现大家都在说的是tensorflow会占满显存的问题（也有可能是我自己没对这个问题表述正确导致没搜索到结果）。

然后思考是不是我的数据有问题。因为虽然数据类似，但是毕竟来源于不同的场景。于是我就回到数据量大的那个场景，从中选了70万条数据训练模型。这个时候观察GPU显存，发现一开始是700多MB，盯着屏幕看了半天，突然等程序运行到1000多步的时候，发现GPU显存一下子跳到了3000多MB。这说明数据量的大小确实会对GPU显存产生影响。

接下来思考这个跳跃点到底是哪个步骤。这样就可以找到为什么GPU显存会一下子升高了。一开始训练一步一步的输出，程序都对，GPU显存也比较低，那为什么中间过程会升高呢？说明一定是到了这一步之后GPU自己存了什么数据，不可能只是模型，因为没有这么大。升高这么多很大可能是存了我们的数据。

怎么验证我们这个猜想呢？我在训练完一轮之后打印了一下，然后看这个打印输出的时候是不是GPU显存升高的时候。事实证明，确实是这样。

结论

当数据量比较小，而GPU的显存足够存储这些数据的时候，在训练过程中会把这些数据存在GPU显存，这样每一轮训练的时候就不用从cpu读数据了，训练速度会快些。而当数据量比较大，而GPU的显存不够时，每次GPU显存中只会存储模型和批数据。

**注：** 1. 以上问题的解决有猜测的成分在，批判着吸收 2. tensorflow实现文本分类（http://www.wildml.com/2015/12/implementing-a-cnn-for-text-classification-in-tensorflow/）

[https_cloud.tencent.com_developer_article_1365920]: https://cloud.tencent.com/developer/article/1365920
[GPU]: https://cloud.tencent.com/product/gpu?from=10680