TensorFlow 设置程序可见GPU与逻辑分区

TensorFlow 设置程序可见GPU(多GPU情况)

import matplotlib as mpl
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import sklearn
import pandas as pd
import os
import sys
import time
import tensorflow as tf

from tensorflow_core.python.keras.api._v2 import keras

print(tf.__version__)
print(sys.version_info)
for module in mpl, np, pd, sklearn, tf, keras:
    print(module.__name__, module.__version__)

# 打印变量所在位置
tf.debugging.set_log_device_placement(True)  

# 获取物理GPU的个数
gpus = tf.config.experimental.list_physical_devices("GPU")  

if len(gpus) >= 1:
    # 设置第几个GPU 当前程序可见
    tf.config.experimental.set_visible_devices(gpus[0], "GPU")
    
print("物理GPU个数:", len(gpus))

# 获取逻辑GPU的个数
logical_gpus = tf.config.experimental.list_logical_devices("GPU") 
print("逻辑GPU个数:", len(logical_gpus))

TensorFlow 设置GPU的 逻辑分区

import matplotlib as mpl
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import sklearn
import pandas as pd
import os
import sys
import time
import tensorflow as tf

from tensorflow_core.python.keras.api._v2 import keras

print(tf.__version__)
print(sys.version_info)
for module in mpl, np, pd, sklearn, tf, keras:
    print(module.__name__, module.__version__)

# 打印变量所在位置
tf.debugging.set_log_device_placement(True)  

# 获取物理GPU的个数
gpus = tf.config.experimental.list_physical_devices("GPU")  

if len(gpus) >= 1:
    # 设置第几个GPU 当前程序可见
    tf.config.experimental.set_visible_devices(gpus[0], "GPU")
    
    # 设置GPU的 逻辑分区
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0], 
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=3072),
         tf.config.experimental.VirtualDeviceConfiguration(memory_limit=3072)])

print("物理GPU个数:", len(gpus))

# 获取逻辑GPU的个数
logical_gpus = tf.config.experimental.list_logical_devices("GPU") 
print("逻辑GPU个数:", len(logical_gpus))

TensorFlow 手动设置处理GPU

import matplotlib as mpl
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import sklearn
import pandas as pd
import os
import sys
import time
import tensorflow as tf

from tensorflow_core.python.keras.api._v2 import keras

print(tf.__version__)
print(sys.version_info)
for module in mpl, np, pd, sklearn, tf, keras:
    print(module.__name__, module.__version__)

# 打印变量所在位置
tf.debugging.set_log_device_placement(True)  

# 自动指定处理设备
tf.config.set_soft_device_placement(True)

# 获取物理GPU的个数
gpus = tf.config.experimental.list_physical_devices("GPU")  
for gpu in gpus:
    # 设置内存自增长方式
    tf.config.experimental.set_memory_growth(gpu, True)  
print("物理GPU个数:", len(gpus))

# 获取逻辑GPU的个数
logical_gpus = tf.config.experimental.list_logical_devices("GPU") 
print("逻辑GPU个数:", len(logical_gpus))

c = []

# 循环遍历当前逻辑GPU
for gpu in logical_gpus:
    print(gpu.name)

    # 手动设置处理GPU
    with tf.device(gpu.name):
        a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
        b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
        
        # 矩阵相乘 并且添加至列表
        c.append(tf.matmul(a, b))

# 手动设置处理GPU
with tf.device("/GPU:0"):
    matmul_sum = tf.add_n(c)

print(matmul_sum)