import torch
import torch_mlu.core.mlu_quantize as mlu_quantize
import torch_mlu
import torch_mlu.core.mlu_model as ct

from nets.facenet import Facenet
from utils.dataloader import LFWDataset
from utils.utils_metrics import test
from utils.utils import preprocess_input
from PIL import Image
import numpy as np

torch.set_grad_enabled(False)

model_path = "model_data/facenet_mobilenet.pth"
backbone = "mobilenet"
model_new_path = "model_data/facenet_mobilenet.pth"
model_int8_path = "model_data/facenet_mobilenet_int8.pth"


def cambricom_save():
    print("###################1################################")
    model = Facenet(mode="predict")
    model.load_state_dict(torch.load(
        model_path, map_location=torch.device('cpu')), strict=False)
    model.eval()
    #torch.save(model, model_new_path, _use_new_zipfile_serialization=False)
    image1 = Image.open("xx.jpg")
    input_data1=torch.from_numpy(np.expand_dims(np.transpose(preprocess_input(np.array(image1, np.float32)), (2, 0, 1)), 0))
    #input_data = torch.randn((1, 3, 160, 160))
    # step 3
    #net_mlu = net.to(ct.mlu_device())
    #input_mlu1 = input_data1.to(ct.mlu_device())
    # step 4
    output1 = model(input_data1).cpu().numpy()
    
    image2 = Image.open("LYK.jpg")
    input_data2=torch.from_numpy(np.expand_dims(np.transpose(preprocess_input(np.array(image2, np.float32)), (2, 0, 1)), 0))
    #input_mlu2 = input_data2.to(ct.mlu_device())
    # step 4
    output2 = model(input_data2).cpu().numpy()
    
    print(output1)
    print(output2)
    l1 = np.linalg.norm(output1 - output2, axis=1)
    print(l1)


def fp_to_int():
    model = Facenet(mode="predict")
    model.load_state_dict(torch.load(
        model_new_path, map_location=torch.device('cpu')), strict=False)
    #model.eval()
    mean = []
    std = []
    # 注意此接口，这里不使用firstconv优化，它的作用是将归一化放到第一层去一起加速做，但是有些模型的前处理是不需要这样做的，具体信息，请参考寒武纪官方文档。
    net_quantization = mlu_quantize.quantize_dynamic_mlu(
        model, {'mean': mean, 'std': std, 'firstconv': False}, dtype='int8', gen_quant=True)
    torch.save(net_quantization.state_dict(), model_int8_path)


def test_int():
    print("###################2################################")
    model = Facenet(mode="predict")
    net = mlu_quantize.quantize_dynamic_mlu(model)

    # step 2
    net.load_state_dict(torch.load(model_int8_path))
    #net.eval()
    # 这里是
    image1 = Image.open("xx.jpg")
    input_data1=torch.from_numpy(np.expand_dims(np.transpose(preprocess_input(np.array(image1, np.float32)), (2, 0, 1)), 0))
    #input_data = torch.randn((1, 3, 160, 160))
    # step 3
    net_mlu = net.to(ct.mlu_device())
    input_mlu1 = input_data1.to(ct.mlu_device())
    # step 4
    output1 = net_mlu(input_mlu1).cpu().numpy()
    
    image2 = Image.open("LYK.jpg")
    input_data2=torch.from_numpy(np.expand_dims(np.transpose(preprocess_input(np.array(image2, np.float32)), (2, 0, 1)), 0))
    input_mlu2 = input_data2.to(ct.mlu_device())
    # step 4
    output2 = net_mlu(input_mlu2).cpu().numpy()
    
    print(output1)
    print(output2)
    l1 = np.linalg.norm(output1 - output2, axis=1)
    print(l1)


def MLU_offline():
    print("###############################################");
    model = Facenet(mode="predict")
    net = mlu_quantize.quantize_dynamic_mlu(model)

    # step 2
    net.load_state_dict(torch.load(model_int8_path))
    # net.eval()
    # 这里是
    input_data = torch.randn((1, 3, 160, 160), dtype=torch.float)
    # step 3
    net_mlu = net.eval().float().to(ct.mlu_device())
    #input_data = input_data.type(torch.HalfTensor)
    input_data = input_data.half()
    input_mlu = input_data.to(ct.mlu_device())

    # 详细查看文档，一般4
    core_number = 1
    ct.set_core_number(core_number)
    ct.set_core_version('MLU220')
    # torch_mlu.core.mlu_model.set_input_format(input_format)
    ct.save_as_cambricon('half_facenet_mobilenet')

    net_trace = torch.jit.trace(net_mlu, input_mlu, check_trace=False)

    net_trace(input_mlu)

    torch_mlu.core.mlu_model.save_as_cambricon("")


cambricom_save()

fp_to_int()
test_int()
#MLU_offline()

模型量化后，输出不对。量化过程中有个warning   ：

量化前后结果对比：