加载两个不同的模型时如何使用model.predict（）函数

Question

我训练了两个CNN，它们都是二进制分类器。我构建了两个分类器脚本，用于加载模型并对图像进行分类，效果良好。在此之前，我还有两个预处理脚本，其中一个只是裁剪并规范化图像，另一个裁剪、删除图像的一部分并规范化图像

在使用分类器脚本之前，我首先让预处理器脚本在图像上运行。现在我想合并所有脚本。因此，首先我的图像应该进行预处理（两个模型使用相同的图像，但集中在不同的部分）。所以我应该有两张图片，一张是模型1，一张是模型2。然后对每个模型使用model.predict()方法并加载相应的图像。但这一次它不起作用。我的一个模型总是预测同一个类，但当我让它在一个脚本中单独运行时，效果很好

上下文：我有可能有划痕或断裂边缘的模具图像。由于两个错误可能同时发生，我训练了两个CNN。一个用于检测划痕，另一个用于检测破损边缘。对于Scratch CNN，我只是裁剪图像，以便图像仅显示模具。对于破碎边缘的CNN，我再次裁剪图像，然后遮住图像的中间，这样CNN就可以只看边缘了。现在我想在同一个脚本中使用这两个模型，这样我就可以整理出坏的模具图像。如果划痕模型看到划痕，则模具的状态为NOK。同样，当断边模型看到断边时，状态设置为NOK。也可能存在划痕和断边，因此这两种状态都设置为NOK。如果没有划痕或断边，两种状态都设置为OK

这导致了4种情况，其中两种状态都可以是OK和NOK。其他两种情况是，每类中有一种为NOK。只有当两者都正常时，模具才应分类为正常；如果至少有一个等级为NOK，则模具应分类为NOK

我当前的代码如下所示：

import os
import time
from shutil import copyfile

import cv2
import numpy as np
from keras.engine.saving import load_model

class ClassifierAgent:
    def __init__(self):
        #load scratch model
        self.scratch_model = load_model("C:\\Users\\but\\PycharmProjects\\OpticalDieInspection\\ScratchModel.h5")
        self.broken_edge_model = load_model("C:\\Users\\but\\PycharmProjects\\OpticalDieInspection\\brokenEdgeModel.h5")

        self.x = 125  # Fixed cropping dimensions for Carsten Problem
        self.y = 54
        self.height = 384
        self.width = 384
        self.dsize = (384, 384)


    def classify(self, input_dir, ok_dir, scratch_dir, broken_edge_dir):
        files = os.listdir(input_dir)
        image_files = [os.path.join(input_dir, f) for f in files if f.endswith(".jpg")]
        total = 0
        good_parts_cnt = 0
        bad_parts_cnt = 0
        for i in range(len(image_files)):
            start_time = time.time()
            total += 1
            img_input = cv2.imread(image_files[i])
            #crop out region of interest (die)
            crop_image = img_input[self.y:self.y + self.height, self.x:self.x + self.width].copy()  # Makes a copy of file
            #normalizing images before processing further
            temp = cv2.normalize(crop_image, None, 0, 255, norm_type=cv2.NORM_MINMAX)
            #get height and width of cropped image
            h, w = crop_image.shape[0:2]

            #copy cropped image
            blacked_image = temp.copy()
            #black out text in the middle of the die
            blacked_image[27:h - 27, 27:w - 27] = 0

            scratch_input_image = np.expand_dims(temp, axis=0)
            #scratch_input_image = scratch_input_image.astype('float32') / 255

            broken_edge_input_image = np.expand_dims(blacked_image, axis=0)
            #broken_edge_input_image = broken_edge_input_image.astype('float32') / 255

            #cv2.imshow("scratch image", scratch_input_image)
            #cv2.imshow("broken edge image", broken_edge_input_image)
            #cv2.waitKey(0)


            output_scratch = self.scratch_model.predict(scratch_input_image)
            output_broken_edge = self.broken_edge_model.predict(broken_edge_input_image)

            if output_scratch[0][0] > output_scratch[0][1]:
                #print(output_scratch[0][0], "Scratch")
                part_state_scratch = "NOK"
            else:
                #print(output_scratch[0][1], 'OK')
                part_state_scratch = "OK"

            if output_broken_edge[0][0] > output_broken_edge[0][1]:
                #print(output_broken_edge[0][0], "BrokenEdge")
                part_state_broken_edge = "NOK"
            else:
                #print(output_broken_edge[0][1], 'OK')
                part_state_broken_edge = "OK"


            if part_state_scratch == "OK" and part_state_broken_edge == "OK":
                ok_dest = os.path.join(ok_dir, os.path.basename(image_files[i]))
                copyfile(image_files[i], ok_dest)
                good_parts_cnt += 1
            elif part_state_scratch == "NOK" and part_state_broken_edge == "OK":
                nok_dest = os.path.join(scratch_dir, os.path.basename(image_files[i]))
                copyfile(image_files[i], nok_dest)
                bad_parts_cnt += 1
            elif part_state_scratch == "OK" and part_state_broken_edge == "NOK":
                nok_dest = os.path.join(broken_edge_dir, os.path.basename(image_files[i]))
                copyfile(image_files[i], nok_dest)
                bad_parts_cnt += 1
            else:
                nok_dest = os.path.join(broken_edge_dir, os.path.basename(image_files[i]))
                copyfile(image_files[i], nok_dest)
                nok_dest = os.path.join(scratch_dir, os.path.basename(image_files[i]))
                copyfile(image_files[i], nok_dest)
                bad_parts_cnt += 1

            end_time = time.time()

            print("Image: ", files[i], "| Scratch: ", part_state_scratch, "; Broken_Edge:  ", part_state_broken_edge, " | ", total, "/", len(image_files))
            print('prediction time:', (end_time - start_time))
            # cv2.imshow("prediction", img_input)

            # cv2.waitKey(0)
            # cv2.destroyAllWindows()


if __name__ == '__main__':
    input_dir = "C:/Users/but/Desktop/Presence/Presence"
    scratch_dir = "C:/Users/but/Desktop/Evaluate/Scratch"
    broken_edge_dir = "C:/Users/but/Desktop/Evaluate/BrokenEdge"
    ok_dir = "C:/Users/but/Desktop/EvaluateBrokenEdge/OK"

    agent = ClassifierAgent()
    agent.classify(input_dir, ok_dir, scratch_dir, broken_edge_dir)

Answer 1

这是在同一台计算机上加载多个模型时发生的常见情况。问题是模型的内存管理。如果您只有一个GPU，那么Keras及其TensorFlow后端将把所有GPU内存分配给一个型号。您可以通过取消分配和分配模型内存来克服这一问题。基本上，从GPU中释放模型并在其中插入新模型。您还可以尝试使用称为会话的

以下是一些有用的链接：

How to release the occupied GPU memory when calling keras model by Apache mod_wsgi django?

Keras: Load multiple models and predict in different threads

https://github.com/keras-team/keras/issues/8538