Пакетирование в tf.Keras 2.1 -> ValueError: Ошибка при проверке ввода

Question

У меня есть набор данных tf, для которого я пытаюсь реализовать пакетирование, чтобы уменьшить нагрузку на память. Модель работает без пакетирования, но я получаю ошибки OOM, следовательно, пробую пакеты.

Набор данных

ds = tf.data.Dataset.from_tensors(( (train_ids, train_masks), labels ))

>>>ds
<TensorDataset shapes: (((22500, 64), (22500, 64)), (22500,)), types: ((tf.int16, tf.int16), tf.int16)>

Пакетирование

BATCH_SIZE = 4
ds = ds.batch(BATCH_SIZE, drop_remainder=True)
TOTAL_BATCHES = math.ceil(len(train_ids) / BATCH_SIZE)
TEST_BATCHES = TOTAL_BATCHES // 10
ds.shuffle(TOTAL_BATCHES)
print("Total Batches : ", str(TOTAL_BATCHES))
print("Test batches : ", str(TEST_BATCHES))

test_data = ds.take(TEST_BATCHES)
val_data = ds.take(TEST_BATCHES)
train_data = ds.skip(TEST_BATCHES*2)

Модель

model = TFAlbertForSequenceClassification.from_pretrained('albert-base-v2', trainable=False)
input_layer = Input(shape=(64, ), dtype=tf.int32)
input_mask_layer = Input(shape=(64, ), dtype=tf.int32)
bert_layer = model([input_layer, input_mask_layer])[0]
flat_layer = Flatten() (bert_layer)
dense_output = Dense(n_classes, activation='softmax') (flat_layer)

model_ = Model(inputs=[input_layer, input_mask_layer], outputs=dense_output)
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3, epsilon=1e-08, clipnorm=1.0)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')

model_.compile(optimizer=optimizer, loss=loss, metrics=[metric])

Сводка модели

Model: "model_17"
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_64 (InputLayer)           [(None, 64)]         0                                            
__________________________________________________________________________________________________
input_65 (InputLayer)           [(None, 64)]         0                                            
__________________________________________________________________________________________________
tf_albert_for_sequence_classifi ((None, 2),)         11685122    input_64[0][0]                   
                                                                 input_65[0][0]                   
__________________________________________________________________________________________________
flatten_18 (Flatten)            (None, 2)            0           tf_albert_for_sequence_classifica
__________________________________________________________________________________________________
dense_17 (Dense)                (None, 45)           135         flatten_18[0][0]                 
==================================================================================================
Total params: 11,685,257
Trainable params: 135
Non-trainable params: 11,685,122

Когда я звоню model_.fit(train_data, validation_data=(val_data), epochs=1, verbose=1) с партией, я получаю ошибка:

        ValueError: Error when checking input: expected input_66 to have 2 dimensions, but got array with shape (None, 22500, 64)

Я знаю, что мне нужно изменить свои входные данные, но я не уверен, как.

Answer 1

Я решил эту проблему, не используя наборы данных tenorflow. Это, очевидно, не идеальное решение, но оно работает на данный момент.

С различными входами и метками в качестве тензоров тф я сделал:

inputs_ids = input_ids.numpy()
input_masks = input_masks.numpy()
labels = labels.numpy()

А затем для подгонки:

model.fit([input_ids, input_masks], labels, epochs=10, batch_size=10)

И теперь пакетная обработка работает и больше не выдает ошибки OOM.