在 tf.Keras 2.1 中進行批處理 -> ValueError: 檢查輸入時出錯

Question

我有一個 tf 數據集，我試圖在其上實施批處理以減少 memory 負載。 model 無需批處理即可工作，但我收到 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

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: "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

我通過不使用 tensorflow 數據集解決了這個問題。 這顯然不是一個完美的修復，但它現在有效。

使用各種輸入和標簽作為 tf 張量，我做了：

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 錯誤。