始终从 tensorflow lite 模型获得 0 的预测

Always getting 0 for prediction from tensorflow lite model

我使用以下代码训练并保存了模型

from tensorflow.keras.preprocessing.image import ImageDataGenerator
import tensorflow as tf
from tensorflow.python.keras.optimizer_v2.rmsprop import RMSprop

train_data_gen = ImageDataGenerator(rescale=1 / 255)
validation_data_gen = ImageDataGenerator(rescale=1 / 255)

# Flow training images in batches of 120 using train_data_gen generator
train_generator = train_data_gen.flow_from_directory(
    'datasets/train/',
    classes=['bad', 'good'],
    target_size=(200, 200),
    batch_size=120,
    class_mode='binary')

validation_generator = validation_data_gen.flow_from_directory(
    'datasets/valid/',
    classes=['bad', 'good'],
    target_size=(200, 200),
    batch_size=19,
    class_mode='binary',
    shuffle=False)

model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(16, (3, 3), activation='relu', input_shape=(200, 200, 3)),
    tf.keras.layers.MaxPooling2D(2, 2),

    tf.keras.layers.Conv2D(32, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2, 2),

    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2, 2),

    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2, 2),

    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2, 2),
    # Flatten the results to feed into a DNN
    tf.keras.layers.Flatten(),
    # 512 neuron hidden layer
    tf.keras.layers.Dense(512, activation='relu'),
    # Only 1 output neuron. It will contain a value from 0-1
    # where 0 for 1 class ('bad') and 1 for the other ('good')
    tf.keras.layers.Dense(1, activation='sigmoid')])

model.compile(loss='binary_crossentropy',
              optimizer=RMSprop(lr=0.001),
              metrics='accuracy')

model.fit(train_generator,
          steps_per_epoch=10,
          epochs=25,
          verbose=1,
          validation_data=validation_generator,
          validation_steps=8)

print("Evaluating the model :")
model.evaluate(validation_generator)

print("Predicting :")

validation_generator.reset()
predictions = model.predict(validation_generator, verbose=1)
print(predictions)

model.save("models/saved")

然后使用

将此模型转换为tflite
import tensorflow as tf


def saved_model_to_tflite(model_path, quantize):
    converter = tf.lite.TFLiteConverter.from_saved_model(model_path)
    model_saving_path = "models/converted/model.tflite"
    if quantize:
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
        model_saving_path = "models/converted/model-quantized.tflite"
    tflite_model = converter.convert()
    with open(model_saving_path, 'wb') as f:
        f.write(tflite_model)

然后使用

测试单个图像的模型
import tensorflow as tf


def run_tflite_model(tflite_file, test_image):

    interpreter = tf.lite.Interpreter(model_path=str(tflite_file))
    interpreter.allocate_tensors()
    print(interpreter.get_input_details())
    input_details = interpreter.get_input_details()[0]
    output_details = interpreter.get_output_details()[0]

    interpreter.set_tensor(input_details["index"], test_image)
    interpreter.invoke()
    output = interpreter.get_tensor(output_details["index"])[0]

    prediction = output.argmax()

    return prediction

main.py

if __name__ == '__main__':


    converted_model = "models/converted/model.tflite"
    bad_image_path = "datasets/experiment/bad/b.png"
    good_image_path = "datasets/experiment/good/g.png"
    img = io.imread(bad_image_path)
    resized = resize(img, (200, 200)).astype('float32')
    test_image = np.expand_dims(resized, axis=0)
    prediction = run_tflite_model(converted_model, test_image)
    print(prediction)

尽管我将图像输入到模型中,但我总是得到 0 的预测。这里有什么问题?

您忘记在将图像传递给 tflite 模型之前对其进行规范化。

resized = resize(img, (200, 200)).astype('float32')
resized = resized / 255.
test_image = np.expand_dims(resized, axis=0)
prediction = run_tflite_model(converted_model, test_image)

编辑:

您正在执行 binary classification task 而不是 multi-class classification task 所以您 do not needoutput array 中的 max value 当作 only produces a single value0 to 1 的范围内。如果值为 greater than or equal to 0.5,则可以将结果解释为 positive example,如果值为 less than 0.5,则可以将结果解释为 negative example

import tensorflow as tf

def run_tflite_model(tflite_file, test_image):
    interpreter = tf.lite.Interpreter(model_path=str(tflite_file))
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    interpreter.set_tensor(input_details[0]["index"], test_image)
    interpreter.invoke()
    predictions = interpreter.get_tensor(output_details[0]["index"])
    return 1 if predictions >= 0.5 else 0 # 1 = good, 0 = bad