tflite 合并模型和条件预测

Question

我有一个问题，希望你能得到一些答案或建议。如何使用 tflite 中前两个模型的预测条件将三个 CNN 模型合并为一个模型？

我首先尝试解释我的实际解决方案和我实施的步骤。

我的想法是，我训练了三个不同的模型来区分 4 个最终类别，前两个模型预测 2 个输出（一个最终类别 (fc) 和一个子类别 (sc)），第三个模型预测 2 个最终类别。

我分别训练了这三个模型，它们具有相同的架构但不同的数据集。这 3 个数据集包含相同的图像，但根据最终类对子类的归属进行不同的组织。与单一模型的 74% 相比，该解决方案为我提供了 89% 的更好测试结果。这就是为什么我想保留这三个模型并进行有条件的组合。

现在我正在使用一个类，该类加载 tflit 模型、加载图像、执行预测并对图像进行分类。如下：

class model1_model2_model3_pipline:
  def __init__(self,model_path):
    self.load_model(model_path)

  def load_model(self,model_path):    
    self.tflite_interpreter = tf.lite.Interpreter(model_path=model_path)
    self.tflite_interpreter.allocate_tensors()
    self.input_details = self.tflite_interpreter.get_input_details()
    self.output_details = self.tflite_interpreter.get_output_details()
    self.d, self.w, self.h, self.ch = self.input_details[0]['shape'] 

  def load_image(self,image_path):
    imarray = cv2.imread(f'{image_path}')
    image = cv2.cvtColor(imarray, cv2.COLOR_BGR2GRAY)
    image = cv2.resize(image, (self.h, self.w))
    image = np.expand_dims(image, 2)
    image = np.array(image,  dtype=self.input_details[0]['dtype'])
    image = np.expand_dims(image, axis=0)
    return image

  def mpredict(self,image):     
    tflite_interpreter = self.tflite_interpreter
    tflite_interpreter.set_tensor(self.input_details[0]['index'], image)
    st = time.time()
    tflite_interpreter.invoke()
    elp = time.time() - st
    return elp

  def classify_image(self, top_k=1):
    output_details = self.output_details[0]
    output = np.squeeze(self.tflite_interpreter.get_tensor(output_details['index']))
    if output_details['dtype'] == np.uint8:
      scale, zero_point = output_details['quantization']
      output = scale * (output - zero_point)
    ordered = np.argpartition(-output, top_k)
    return [(i, output[i]) for i in ordered[:top_k]]

  def predict_single_image(self,image_path):
    img = self.load_image2(image_path)
    self.mpredict(img)
    out = self.classify_image(top_k=1)
    return out

现在我必须创建该类的三个实例，并为每个模型预测定义条件来执行预测。

# fc: final-class, sc: sub-class
MAP_CHARACTERS_model1 = {
  0: 'fc1',
  1: 'sc1',
  }
MAP_CHARACTERS_model2 = {
  0: 'fc2',
  1: 'sc2',
  }

MAP_CHARACTERS_model3 = {
  0: 'fc3',
  1: 'fc4',
  }

x_model1 = model1_model2_model3_pipline(model1_pth) #pth for path
x_model2 = model1_model2_model3_pipline(model2_pth)
x_model3 = model1_model2_model3_pipline(model3_pth)

out_model1 = x_model1.predict_single_image(im_pth)
predicted = MAP_CHARACTERS_model1[out_model1[0][0]]

if predicted == 'sc1':
  out_model2 = x_model2.predict_single_image(im_pth)
  predicted = MAP_CHARACTERS_model2[out_model2[0][0]]
  if predicted =='sc2':
    out_model3 = x_model3.predict_single_image(im_pth)
    predicted = MAP_CHARACTERS_model3[out_model3[0][0]]

我想将这三个模型合并为一个模型，并且只需一步即可执行预测。有没有解决方案可以在 tflite 上执行此过程？

Answer 1

简单/推荐的方法是在 TF 中添加额外的调节。这将为您提供 TF 中的一个大模型，可以满足您的需求，然后将此合并模型转换为 TFLite。