AssertionError: batch_size должен делиться на количество используемых ядер TPU (1 против 8) при использовании функции прогнозирования

Question

Некоторые детали для контекста:

1. Работа в Google Colab с использованием TPU.
2. Модель подходит без проблем
3. Проблемы при попытке использовать функцию прогнозирования

Вот код, который я использую для обучения:

tpu_model.fit(x, y,
          batch_size=128,
          epochs=60)

Вот код, который я использую для предсказания:

def generate_output():
    generated = ''
    #sentence = text[start_index: start_index + Tx]
    #sentence = '0'*Tx
    usr_input = input("Write the beginning of your poem, the Shakespeare machine will complete it. Your input is: ")
    # zero pad the sentence to Tx characters.
    sentence = ('{0:0>' + str(maxlen) + '}').format(usr_input).lower()
    generated += usr_input 

    sys.stdout.write("\n\nHere is your poem: \n\n") 
    sys.stdout.write(usr_input)
    for i in range(400):

        x_pred = np.zeros((1, maxlen, len(chars)))

        for t, char in enumerate(sentence):
            if char != '0':
                x_pred[0, t, char_indices[char]] = 1.

        --> preds = tpu_model.predict(x_pred, batch_size = 128 ,workers = 8,verbose=0)[0]
        next_index = sample(preds, temperature = 1.0)
        next_char = indices_char[next_index]

        generated += next_char
        sentence = sentence[1:] + next_char

        sys.stdout.write(next_char)
        sys.stdout.flush()

        if next_char == '\n':
            continue

А вот и ошибка (добавлена ​​стрелка выше, чтобы вы знали местонахождение ошибки:

AssertionError: batch_size must be divisible by the number of TPU cores in use (1 vs 8)

Это не имеет смысла для меня, так как размер партии, который я использовал во время обучения, делится на 8 И размер партии, который я передаю в моей функции прогнозирования, делится на 8.

Я не уверен, в чем проблема и как ее решить. Любая помощь приветствуется.

Answer 1

с ошибкой:

AssertionError: batch_size must be divisible by the number of TPU cores in use (1 vs 8)

Похоже, вы используете batch_size 1, который может быть выведен из первого измерения ваших входных данных:

x_pred = np.zeros((1, maxlen, len(chars)))

Я думаю, вы можете изменить его на:

x_pred = np.zeros((8, maxlen, len(chars)))

, чтобы размер пакета стал равным 8, что соответствует количеству используемых ядер TPU.

Или вы можете сохранить текущий размер пакета batch_size, но использовать 1 ядро ​​TPU.