Search code examples
tensorflowtransfer-learningvgg-net

I cannot train my model using transfer learning (VGG16) on my GPU

I am trying to use transfer learning to train my model to detect diseases in rice on images of the plant. I attempted using VGG16, but I could not get it to train with my GPU. I have an NVIDIA GeForce MX150.

Below is the code that I used to fit the model:

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)
    )
Solution

  • Tensorflow GPU support requires a few dependencies. Please see https://www.tensorflow.org/install/gpu

    Then, try tf.test.is_gpu_available() - if this is True, then your GPU is being used for training.

    On a single GPU, you should not need to use with to train on GPU. For me to help more, please provide any logs or errors.