【环境搭建】cambricon测试

# -*- coding: utf-8 -*-
import torch_mlu
import torch_mlu.core.mlu_model as ct
import torch_mlu.core.mlu_quantize as mlu_quantize
import torch
import numpy as np
import torchvision.transforms as transforms
from torchvision.models.resnet import resnet50
from PIL import Image

torch.set_grad_enabled(False)

def get_torch_image(path, resize=256, crop=224, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
    transform = transforms.Compose([
        transforms.Resize(resize),
        transforms.CenterCrop(crop),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ])
    img_mat = Image.open(path)
    img_mat = transform(img_mat)
    im_tensor = torch.unsqueeze(img_mat, 0)
    return im_tensor


if __name__ == '__main__':
    net = resnet50(pretrained=True)
    net.eval().float()

    im_tensor = get_torch_image("./imagenet/3.jpg")
    outputs = net(im_tensor).detach().numpy().reshape(-1)
    classify = np.argmax(outputs)
    print(classify)

    # combricon quant
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    qconfig = {"iteration": 1, "use_avg": False, "data_scale": 1.0, "firstconv": True, "mean": mean, "std": std,
               "per_channel": False}
    quantized_net = torch_mlu.core.mlu_quantize.quantize_dynamic_mlu(net.eval().float(), qconfig_spec=qconfig,
                                                                     dtype="int8", gen_quant=True)
    quantized_net = quantized_net.eval().float()  # model -> model.eval
    im_tensor = get_torch_image("./imagenet/3.jpg")
    outputs = quantized_net(im_tensor).detach().numpy().reshape(-1)
    classify = np.argmax(outputs)
    print(classify)

    # combricon quant save combricon
    mean = [0, 0, 0]
    std = [1 / 255, 1 / 255, 1 / 255]
    torch.save(quantized_net.state_dict(), "resnet50_quant.pth")
    state_dict = torch.load("resnet50_quant.pth")
    quantized_net = torch_mlu.core.mlu_quantize.quantize_dynamic_mlu(net)
    quantized_net.load_state_dict(state_dict)
    quantized_net_mlu = quantized_net.to(ct.mlu_device())

    ct.set_core_number(1)
    ct.set_core_version("MLU220")
    ct.save_as_cambricon("resnet50_1batch_1core") # 设置保存路径

    im_tensor = get_torch_image("imagenet/3.jpg", mean=mean, std=std)
    im_tensor = im_tensor.to(ct.mlu_device())
    net_traced = torch.jit.trace(quantized_net_mlu, im_tensor, check_trace=False)

    outputs = net_traced(im_tensor).cpu().detach().numpy().reshape(-1)
    classify = np.argmax(outputs)
    print(classify)