TensorFlow 2.0 在单 GPU 上训练模型

Question

我想在单个 NVIDIA 显卡（RTX 2080 super）上训练一个顺序 tensorflow（版本 2.3.0）模型。我正在使用以下代码 sn-p 来构建和训练模型。但是，每次运行此代码时，我都看不到任何 GPU 利用率。任何建议如何修改我的代码以便我可以在 1 个 GPU 上运行它？

strategy = tf.distribute.OneDeviceStrategy(device="/GPU:0")
with strategy.scope():
    num_classes=len(pd.unique(cats.No))
    model = BuildModel((image_height, image_width, 3), num_classes)
    model.summary()
    model=train_model(model,valid_generator,train_generator,EPOCHS,BATCH_SIZE)

Answer 1

运行以下代码以查看 tensorflow 是否检测到您的 GPU。

import tensorflow as tf
from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())
print(tf.__version__)
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
tf.test.is_gpu_available()
!python --version