简体   繁体   中英

Issue using keras with tensorflow backend on GPU

I'm having a problem using Keras with Tensorflow backend. I'm trying to train a keras.applications.VGG16 model without the top layers, replacing them with two dense layers (1024 and 2 nodes) for a binary classification problem. I'm training the model in a NVidia Geforce GTX 1050 with 2GB of dedicated VRAM.

The issue here is that when I start the training, an error occures:

NotFoundError (see above for traceback): No registered 'Snapshot' OpKernel for GPU devices compatible with node training/SGD/gradients/loss/dense_two_loss/Mean_2_grad/floordiv = Snapshot[T=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](training/SGD/gradients/loss/dense_two_loss/Mean_2_grad/Prod/_289)
        .  Registered:  device='CPU'; T in [DT_INT64]
  device='CPU'; T in [DT_INT32]
  device='CPU'; T in [DT_UINT16]
  device='CPU'; T in [DT_INT16]
  device='CPU'; T in [DT_UINT8]
  device='CPU'; T in [DT_INT8]
  device='CPU'; T in [DT_HALF]
  device='CPU'; T in [DT_BFLOAT16]
  device='CPU'; T in [DT_FLOAT]
  device='CPU'; T in [DT_DOUBLE]
  device='CPU'; T in [DT_COMPLEX64]
  device='CPU'; T in [DT_COMPLEX128]
  device='CPU'; T in [DT_BOOL]

[[Node: training/SGD/gradients/loss/dense_two_loss/Mean_2_grad/floordiv = Snapshot[T=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](training/SGD/gradients/loss/dense_two_loss/Mean_2_grad/Prod/_289)]]

What could be the problem here? My code is the following:

from keras.models import Model
from keras.optimizers import SGD
from keras.layers import Dense, GlobalAveragePooling2D
from keras.applications.vgg16 import VGG16
from keras.utils import to_categorical

base_model = VGG16(weights='imagenet', include_top=False)
x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dense(1024, activation='relu', name="dense_one")(x)
predictions = Dense(2, activation='softmax', name="dense_two")(x)

model = Model(inputs=base_model.input, outputs=predictions)

model.compile(optimizer=SGD(lr=0.001, decay=1e-6, momentum=0.5, nesterov=True),
    loss='categorical_crossentropy',
    metrics=['accuracy','mse'])

model.fit(x_train, to_categorical(y_train, num_classes=2), batch_size=1, epochs=150, verbose=2, validation_split=0.1, shuffle=True)

Problem solved. I had installed the 1.7 nightly version of tensorflow and forgot to remove it. I installed 1.6 version and it worked :)

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM