首页 > 解决方案 > Tensorflow:批量分类图像

问题描述

我已经按照这个 TensorFlow 教程使用迁移学习方法对图像进行分类。使用在预训练的 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

标签: tensorflowbatch-processing

解决方案


替换此代码:

# 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()参考


推荐阅读