我不能在GPU上使用传输学习(VGG16)来训练我的模型。

Question

我正在尝试利用转移学习来训练我的模型，在植物图像上检测水稻中的疾病。我尝试使用VGG16，但我无法让它与我的GPU训练。我有一个NVIDIA GeForce MX150。

下面是我用来拟合模型的代码：

import tensorflow as tf
from tensorflow.python.client import device_lib

print(device_lib.list_local_devices())

with tf.device('/device:GPU:1'):
    # fit the model
    r = model.fit(
        training_set,
        validation_data=test_set,
        epochs=20,
        steps_per_epoch=len(training_set),
        validation_steps=len(test_set)
    )

Answer 1

Tensorflow GPU支持需要一些依赖项。请参阅https://www.tensorflow.org/install/gpu

然后，尝试tf.test.is_gpu_available() -如果这是True，那么您的GPU将被用于培训。

在单个GPU上，您不需要使用with来训练GPU。对于我的帮助更多，请提供任何日志或错误。