量化TFLite格式MobilenetV2分类模型Edge TPU部署结果异常问询

目标是将PyTorch模型转换为量化TFLite模型，用于Edge TPU上的推理。

• 已完成较为复杂的深度估计模型从PyTorch到TFLite格式的转换，可在Edge TPU上正常运行，但由于部分算子不支持，推理速度较慢（超过800ms），算子运行统计如下：

Number of operations that will run on Edge TPU: 87
Number of operations that will run on CPU: 47

[深度估计模型结构示意图]

• 为得到可完全在TPU上运行的模型，尝试转换了结构最简的MobilenetV2分类模型，但运行该量化模型时得到的结果偏差极大，具体对比如下：

该问题是float32转uint8的量化过程导致的，还是操作存在错误？如果是量化导致的，该如何缓解该问题？已知Coral官方提供的同模型分类示例可正常运行。

转换链路为：PyTorch -> ONNX -> OpenVINO -> TensorFlow -> TensorFlowLite
自行实现了PyTorch转ONNX、TensorFlow(pd)转TFLite的代码，其余转换步骤使用OpenVINO mo.py脚本和openvino2tensorflow工具完成，用于解决PyTorch和TensorFlow之间的NCHW、NHWC格式不匹配问题。

深度估计模型、MobilenetV2分类模型、ImageNet标签文件、测试图片均可在对应公开开源仓库获取。

以下代码无需Edge TPU即可运行，但需要安装Google Coral相关依赖库。如果使用(2.0, 76.0)这类均值和标准差参数，单张测试狗图的识别结果会正常，但其他图片的识别仍然存在同样的偏差问题。

TFLite模型测试代码 #

import numpy as np
from PIL import Image
from pycoral.adapters import classify
from pycoral.adapters import common
from pycoral.utils.dataset import read_label_file
from torchvision import transforms

from tensorflow.lite.python.interpreter import Interpreter


def cropPIL(image, new_width, new_height):
    width, height = image.size

    left = (width - new_width)/2
    top = (height - new_height)/2
    right = (width + new_width)/2
    bottom = (height + new_height)/2

    return image.crop((left, top, right, bottom))

def softmax(x):
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum()

def classify_img(image_dir, lables_dir, model_dir, mean, std):
    # 加载标签和模型
    labels = read_label_file(lables_dir)
    interpreter = Interpreter(model_path=model_dir)
    interpreter.allocate_tensors()
    
    # 加载并 resize 图片
    size = (256, 256)
    image = Image.open(image_dir).convert('RGB')
    image = image.resize(((int)(size[0]*image.width/image.height), size[1]), Image.ANTIALIAS)
    image = cropPIL(image, 224, 224)
    image = np.asarray(image)

    # 输入图片归一化
    params = common.input_details(interpreter, 'quantization_parameters')
    scale = params['scales']
    zero_point = params['zero_points']

    normalized_input = (image - mean) / (std * scale) + zero_point
    np.clip(normalized_input, 0, 255, out=normalized_input)

    # 输入喂入模型
    common.set_input(interpreter, normalized_input.astype(np.uint8))
    
    # 执行推理
    interpreter.invoke()

    # 获取输出张量并执行softmax
    output_details = interpreter.get_output_details()[0]
    output_data = interpreter.tensor(output_details['index'])().flatten()
    scores = softmax(output_data.astype(float))

    # 获取Top5分类结果
    classes = classify.get_classes_from_scores(scores, 5, 0.0)

    print('-------结果输出--------')
    for c in classes:
       print('%s: %f' % (labels.get(c.id, c.id), c.score))


image_dir  = 'data/dog.jpg'
lables_dir = 'data/imagenet_classes.txt'
model_dir  = 'models/mobilenetv2.tflite'

classify_img(image_dir, lables_dir, model_dir, 114.0, 57.0)

要在Google Colab上运行PyTorch模型，需要将
model = torch.hub.load('pytorch/vision:v0.9.0', 'mobilenet_v2', pretrained=True)
替换为
model = torchvision.models.mobilenet_v2(pretrained=True)
才可正常运行。

PyTorch原生模型测试代码 #

import torch
from PIL import Image
from torchvision import transforms
import torchvision

import numpy as np
import matplotlib.pyplot as plt
    
def inference(model, input_image, lables_dir):
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    input_tensor = preprocess(input_image)
    input_batch = input_tensor.unsqueeze(0)

    # 如有GPU可将输入和模型移至GPU加速
    if torch.cuda.is_available():
        input_batch = input_batch.to('cuda')
        model.to('cuda')

    with torch.no_grad():
        output = model(input_batch)

    probabilities = torch.nn.functional.softmax(output[0], dim=0)

    # 读取类别标签
    with open(lables_dir, "r") as f:
        categories = [s.strip() for s in f.readlines()]

    # 获取Top5分类结果
    top5_prob, top5_catid = torch.topk(probabilities, 5)
    result = {}
    for i in range(top5_prob.size(0)):
        result[categories[top5_catid[i]]] = top5_prob[i].item()
    return result

def classify(image_dir, lables_dir):
    model = torchvision.models.mobilenet_v2(pretrained=True)
    model.eval()

    im = Image.open(image_dir)
    results = inference(model, im, lables_dir)
    for result in results:
        print(f'{result}: {round(results[result], 5)}')


classify('data/dog.jpg', 'data/imagenet_classes.txt')

内容的提问来源于stack exchange，提问作者Aaron Zettler