gpus = tf.config.experimental.list_physical_devices(device_type='GPU') tf.config.experimental.set_memory_growth(gpus[0], True)
时间: 2023-12-06 09:44:11 浏览: 162
人工智能-tf-测试gpu是否可用
这段代码的作用是列出所有可用的物理GPU设备,并设置GPU内存动态增长模式,以便在需要更多GPU内存时自动分配。
首先,`tf.config.experimental.list_physical_devices(device_type='GPU')`会列出所有可用的物理GPU设备。
然后,`tf.config.experimental.set_memory_growth(gpus[0], True)`会将第一个GPU设备的内存动态增长模式设置为 True。这意味着当需要更多GPU内存时,TensorFlow会自动分配额外的内存,而不是在程序开始运行时分配所有内存。这样可以避免内存浪费,并确保程序在需要更多内存时不会崩溃。
阅读全文