tensorflow-gpu设置显存占用情况（Memory-Usage）

原创于 2023-10-26 14:14:24 发布 · 718 阅读

4 ·

CC 4.0 BY-SA版权

文章标签：

#tensorflow #人工智能 #python #深度学习

文章讲述了在服务器上运行深度学习代码时，由于数据量小导致内存占用过高的问题，通过导入tensorflow并使用`tf.config`设置虚拟设备配置，限制GPU内存为2048MB来解决内存占用问题。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

##起因

每次在服务器跑深度学习代码，我的数据量很小。但是内存直接占满。如下图

可以设置内存占用情况。

解决办法

### 在前三行加上代码：

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
tf.config.experimental.set_virtual_device_configuration(gpus[0],[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)])

结果

解决内存占用问题