重塑BatchDataset训练模型的输入-Tensorflow

在将tensorflow模型转换为BatchDataset后,我在图像上对其进行了训练

  IMG_size = 224
  INPUT_SHAPE = [None, IMG_size, IMG_size, 3] # 4D input

  model.fit(x=train_data,
            epochs=EPOCHES,
            validation_data=test_data,
            validation_freq=1, # check validation metrics every epoch
            callbacks=[tensorboard, early_stopping])


     model.compile(
      loss=tf.keras.losses.CategoricalCrossentropy(), 
      optimizer=tf.keras.optimizers.Adam(), 
      metrics=["accuracy"] 
    )


  model.build(INPUT_SHAPE)

“列车数据”类型为:tensorflow。python数据行动。数据集操作。BatchDataset。
我想在单个numpy数组或张量常量上运行我的模型,但它将是3D输入矩阵,而不是4D作为输入“TensorShape([224,224,3]),我如何重塑它?