tf.argmax не работает при использовании на GPU (Tensorflow) - PullRequest
0 голосов
/ 24 января 2019

У меня проблема с использованием tf.argmax на GPU. Проблема заключается в следующем

следующий код

for i in range(10): # Loop for the Epochs
  print ("\nEpoch:", i)

  for (batch, (images, labels)) in enumerate(dataset.take(60000)): # Loop for the mini-batches
    if batch % 100 == 0:
      #print('batches processed', batch)
      print('.', end='')
    labels = tf.cast(labels, dtype = tf.int64)

    with tf.device('/gpu:0'):
      with tf.GradientTape() as tape:
        logits = mnist_model(images, training=True)
        #print(logits)
        i64 = tf.constant(1, dtype=tf.int64)
        tgmax = tf.argmax(labels, axis = i64)
        loss_value = tf.losses.sparse_softmax_cross_entropy(tgmax, logits)

        loss_history.append(loss_value.numpy())
        grads = tape.gradient(loss_value, mnist_model.variables)
        optimizer.apply_gradients(zip(grads, mnist_model.variables),
                                    global_step=tf.train.get_or_create_global_step())

работает на CPU, но не (как написано) на GPU. Это простой тест с использованием активного выполнения в наборе данных MNIST. Я думаю, что ошибка связана с типами данных, но я не совсем уверен. Я получаю ошибку

NotFoundError: No registered 'ArgMax' OpKernel for GPU devices compatible with node {{node ArgMax}} = ArgMax[T=DT_INT64, Tidx=DT_INT64, output_type=DT_INT64](dummy_input, dummy_input)
     (OpKernel was found, but attributes didn't match)
    .  Registered:  device='XLA_GPU'; output_type in [DT_INT32, DT_INT64]; T in [DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT8, DT_COMPLEX64, DT_INT64, DT_QINT8, DT_QUINT8, DT_QINT32, DT_BFLOAT16, DT_HALF, DT_UINT32, DT_UINT64]; Tidx in [DT_INT32, DT_INT64]
  device='XLA_GPU_JIT'; output_type in [DT_INT32, DT_INT64]; T in [DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT8, DT_COMPLEX64, DT_INT64, DT_QINT8, DT_QUINT8, DT_QINT32, DT_BFLOAT16, DT_HALF, DT_UINT32, DT_UINT64]; Tidx in [DT_INT32, DT_INT64]
  device='GPU'; T in [DT_DOUBLE]; output_type in [DT_INT32]; Tidx in [DT_INT32]

У кого-нибудь есть простое решение? Заранее спасибо, Умберто

1 Ответ

0 голосов
/ 25 января 2019

Выглядит как ошибка, как упомянуто здесь . Пожалуйста, сообщите о проблеме в tenorflow на github.

Добро пожаловать на сайт PullRequest, где вы можете задавать вопросы и получать ответы от других членов сообщества.
...