Tensorflow：批量分类图像

Question

我已关注this TensorFlow tutorial，使用迁移学习方法对图像进行分类。使用在预训练的 MobileNet V2 模型之上添加的近 16,000 个手动分类图像（大约 40/60 分割为 1/0），我的模型在保留测试集上实现了 96% 的准确率。然后我保存了生成的模型。

接下来，我想使用这个经过训练的模型对新图像进行分类。为此，我以下述方式调整了教程代码的一部分（最后是 #Retrieve abatch of the test set）。该代码有效，但是，它只处理一批 32 张图像，仅此而已（源文件夹中有数百张图像）。我在这里想念什么？请指教。

# Import libraries
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import preprocessing
from tensorflow.keras.preprocessing import image_dataset_from_directory
import matplotlib.pyplot as plt
import numpy as np
import os

# Load saved model
model = tf.keras.models.load_model('/model')

# Re-compile model
base_learning_rate = 0.0001
model.compile(optimizer=tf.keras.optimizers.Adam(lr=base_learning_rate),
              loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
              metrics=['accuracy'])
# Define paths
PATH = 'Data/'             
new_dir = os.path.join(PATH, 'New_images') # New_images must contain at least one class (sub-folder)

IMG_SIZE = (640, 640)
BATCH_SIZE = 32

new_dataset = image_dataset_from_directory(new_dir, shuffle=True, batch_size=BATCH_SIZE, image_size=IMG_SIZE)

# Retrieve a batch of images from the test set
image_batch, label_batch = new_dataset.as_numpy_iterator().next()
predictions = model.predict_on_batch(image_batch).flatten()

# Apply a sigmoid since our model returns logits
predictions = tf.nn.sigmoid(predictions)
predictions = tf.where(predictions < 0.5, 0, 1)

print('Predictions:\n', predictions.numpy())
len(new_dataset) # equals 25, i.e., there are 25 batches

Answer 1

替换此代码：

# Retrieve a batch of images from the test set
image_batch, label_batch = new_dataset.as_numpy_iterator().next()
predictions = model.predict_on_batch(image_batch).flatten()

用这个：

predictions = model.predict(new_dataset,batch_size=BATCH_SIZE).flatten()

tf.data.Dataset 对象可以直接传递给方法predict()。 Reference