35、基于shufflenetv2分类模型修改和onnx模型转化，以及ncnn模型部署

基本思想：使用简单的shufflenetv2进行分类，然后转成onnx，部署ncnn上

一、代码fork别的大佬的，源代码：​​GitHub - Randl/ShuffleNetV2-pytorch: Implementation of ShuffleNetV2 for pytorch​​

代码在另一个大佬基础上做了修改，我fork过来了：​​https://github.com/sxj731533730/shufflenetv2_demo​​

链接：https://pan.baidu.com/s/1-33wTFKlzAcNWgn9DYFphA 
提取码：1j6s

预训练模型和训练之后的模型

链接：https://pan.baidu.com/s/1ZO0T0z-DKdWooN32rhJj3w 
提取码：28x5

测试图片

测试目录结果

首先先训练模型，生成pt文件，测试一下

"C:\Program Files\Python36\python.exe" G:/sxj731533730/sxj731533730/csdn/Test7_shufflenet/predict.py
tensor([-0.9376, -1.6345,  0.4400, -0.8859, -1.2804,  0.7720, -1.1510,  3.4949,
         2.9412, -1.6998])
tensor([0.0068, 0.0034, 0.0271, 0.0072, 0.0048, 0.0377, 0.0055, 0.5742, 0.3301,
        0.0032])
class: racing car   prob: 0.574

Process finished with exit code 0

然后进行模型onnx转化

import torch
# import trochvision
import torch.utils.data
import argparse
import onnxruntime
from model import shufflenet_v2_x1_0
import os
import cv2
import numpy as np
from torch.autograd import Variable
from onnxruntime.datasets import get_example
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt

def main(args):
    # print("the version of torch is {}".format(torch.__version__))
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    dummy_input=getInput(args.img_size)#获得网络的输入
    # 加载模型
    model = torch.load(args.model_path)
    model.to('cpu')
    model.eval（)
    pre=model(dummy_input)
    print("the pre:{}".format(pre))
    #保存onnx模型
    torch2onnx(args,model,dummy_input)

def getInput(img_size):
    data_transform = transforms.Compose(
        [#transforms.Resize(256),
         #transforms.CenterCrop(224),
         transforms.ToTensor(),#数据值从[0,255]范围转为[0,1]，相当于除以255操作
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

    # load image
    img_path = r"data_set\test\0abad4440c4415b9707f11151762526c.jpg"
    assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
    img = Image.open(img_path)

    # [N, C, H, W]
    img = data_transform(img)
    # expand batch dimension
    dummy_input = torch.unsqueeze(img, dim=0)
    return dummy_input


def torch2onnx(args,model,dummy_input):
    input_names = ['input']#模型输入的name
    output_names = ['output']#模型输出的name
    # return
    torch_out = torch.onnx._export(model, dummy_input, os.path.join(args.save_model_path,"model-299.onnx"),
                                   verbose=True, input_names=input_names, output_names=output_names)
    # test onnx model
    example_model = get_example(os.path.join(args.save_model_path,"model-299.onnx"))
    session = onnxruntime.InferenceSession(example_model)
    # get the name of the first input of the model
    input_name = session.get_inputs()[0].name
    print('onnx Input Name:', input_name)
    result = session.run([], {input_name: dummy_input.data.numpy()})
    # np.testing.assert_almost_equal(
    #     torch_out.data.cpu().numpy(), result[0], decimal=3)
    print("the result is {}".format(result[0]))
    #结果对比--有点精度上的损失
    # pytorch tensor([[ 5.8738, -5.4470]], device='cuda:0')
    # onnx [ 5.6525207 -5.2962923]

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="PyTorch model to onnx and ncnn")
    parser.add_argument('--model_path', type=str, default=r"G:\sxj731533730\sxj731533730\csdn\Test7_shufflenet\weights\model-299.pth",
                        help="For training from one model_file")
    parser.add_argument('--save_model_path', type=str, default=r"G:\sxj731533730\sxj731533730\csdn\Test7_shufflenet\weights",
                        help="For training from one model_file")
    parser.add_argument('--onnx_model_path', type=str, default=r"G:\sxj731533730\sxj731533730\csdn\Test7_shufflenet\weights",
                        help="For training from one model_file")
    parser.add_argument('--img_size', type=int, default=256,
                        help="the image size of model input")
    args = parser.parse_args()
    main(args)

另一个转化代码，只是读取图片方式不同

import torch
# import trochvision
import torch.utils.data
import argparse
import onnxruntime
from model import shufflenet_v2_x1_0
import os
import cv2
import numpy as np
from torch.autograd import Variable
from onnxruntime.datasets import get_example
from torchvision import transforms
from PIL import Image
def main(args):
    # print("the version of torch is {}".format(torch.__version__))
    dummy_input=getInput(args.img_size)#获得网络的输入
    # 加载模型
    model = torch.load(args.model_path)
    model.to('cpu')
    model.eval（)
    pre=model(dummy_input)
    print("the pre:{}".format(pre))
    #保存onnx模型
    torch2onnx(args,model,dummy_input)

def getInput(img_size):

    input = cv2.imread(r"data_set\test\0abad4440c4415b9707f11151762526c.jpg")
    #input = cv2.resize(input, (img_size, img_size))  # hwc bgr
    input = cv2.cvtColor(input, cv2.COLOR_BGR2RGB)  # hwc rgb
    # [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    input = np.transpose(input, (2, 0, 1)).astype(np.float32)  # chw rgb
    # input=input/255.0
    print("befor the input[0,0,0]:{}".format(input[0, 0, 0]))
    print("the size of input[0,...] is {}".format(input[0, ...].shape))
    print("the size of input[1,...] is {}".format(input[1, ...].shape))
    print("the size of input[2,...] is {}".format(input[2, ...].shape))
    input[0, ...] = ((input[0, ...]/255.0)-0.485)/0.229
    input[1, ...] = ((input[1, ...]/255.0)-0.456)/0.224
    input[2, ...] = ((input[2, ...]/255.0)-0.406)/0.225
    print("after input[0,0,0]:{}".format(input[0, 0, 0]))

    now_image1 = Variable(torch.from_numpy(input))
    dummy_input = now_image1.unsqueeze(0)

    return dummy_input


def torch2onnx(args,model,dummy_input):
    input_names = ['input']#模型输入的name
    output_names = ['output']#模型输出的name
    # return
    torch_out = torch.onnx._export(model, dummy_input, os.path.join(args.save_model_path,"model-299_0.onnx"),
                                   verbose=True, input_names=input_names, output_names=output_names)
    # test onnx model
    example_model = get_example(os.path.join(args.save_model_path,"model-299_0.onnx"))
    session = onnxruntime.InferenceSession(example_model)
    # get the name of the first input of the model
    input_name = session.get_inputs()[0].name
    print('Input Name:', input_name)
    result = session.run([], {input_name: dummy_input.data.numpy()})
    # np.testing.assert_almost_equal(
    #     torch_out.data.cpu().numpy(), result[0], decimal=3)
    print("the result is {}".format(result[0]))
    #结果对比--有点精度上的损失
    # pytorch tensor([[ 5.8738, -5.4470]], device='cuda:0')
    # onnx [ 5.6525207 -5.2962923]

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="PyTorch model to onnx and ncnn")
    parser.add_argument('--model_path', type=str,
                        default=r"G:\sxj731533730\sxj731533730\csdn\Test7_shufflenet\weights\model-299.pth",
                        help="For training from one model_file")
    parser.add_argument('--save_model_path', type=str,
                        default=r"G:\sxj731533730\sxj731533730\csdn\Test7_shufflenet\weights",
                        help="For training from one model_file")
    parser.add_argument('--onnx_model_path', type=str,
                        default=r"G:\sxj731533730\sxj731533730\csdn\Test7_shufflenet\weights",
                        help="For training from one model_file")
    parser.add_argument('--img_size', type=int, default=256,
                        help="the image size of model input")
    args = parser.parse_args()
    main(args)

测试pt 结果为

the pre:tensor([[-0.9376, -1.6345,  0.4400, -0.8859, -1.2804,  0.7720, -1.1510,  3.4949,
          2.9412, -1.6998]], grad_fn=<AddmmBackward0>)

转成onnx输出结果

Input Name: input
the result is [[-0.9375885 -1.6344597  0.4399762 -0.8858576 -1.2804033  0.7720084
  -1.1509724  3.4949338  2.9412236 -1.6998374]]

这里主要有这么一个问题：

因为你要转模型进行模型间的数据对比，最好把predict.py中的transform中resize注释掉，这主要是因为模型训练使用是PIL读取的图片其中使用transform中resize和opencv中的resize对图片处理存在差异，所以使用官方代码predict,py预测（经过transform中的resize）出的结果和使用上述两个例程的代码测试pt和onnx结果是存在差异的，如果你将predict.py的transform中resize函数注释掉，转模型中的resize函数和即将运行ncnn推理函数reisze都注释掉，将保证这三方推理结果是一致的~

然后使用ncnn的onnx2ncnn转一下模型到param和bin （使用我之前用MinGW编译的onnx2ncnn可执行程序）

D:\ncnn\buildMinGW\install\bin>onnx2ncnn.exe G:\sxj731533730\sxj731533730\csdn\Test7_shufflenet\weights\model-299.onnx G:\sxj731533730\sxj731533730\csdn\Test7_shufflenet\weights\model-299.param G:\sxj731533730\sxj731533730\csdn\Test7_shufflenet\weights\model-299.bin

D:\ncnn\buildMinGW\install\bin>

模型

ncnn的测试代码

#include <stdio.h>
#include <algorithm>
#include <vector>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>

#include "platform.h"
#include "net.h"
#if NCNN_VULKAN
#include "gpu.h"
#endif // NCNN_VULKAN

static int detect_squeezenet(const cv::Mat& bgr, std::vector<float>& cls_scores)
{
    ncnn::Net shufflenetv2;




    shufflenetv2.load_param("/home/ubuntu/CLionProjects/untitled2/model-299_0.param");
    shufflenetv2.load_model("/home/ubuntu/CLionProjects/untitled2/model-299_0.bin");

    ncnn::Mat in = ncnn::Mat::from_pixels(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, bgr.cols, bgr.rows);
    //按照pytorch样例进行标准化
    const float mean_vals[3] = {0.485f, 0.456f, 0.406f};
    const float std_vals[3] = {1/0.229f, 1/0.224f, 1/0.225f};
    const float norm_255[3] = {1/255.0f, 1/255.0f, 1/255.0f};
    // in.substract_mean_normalize(mean_vals, std_vals);
    in.substract_mean_normalize(0, norm_255);
    in.substract_mean_normalize(mean_vals, std_vals);
    fprintf(stderr, "input shape: %d %d %d %d\n", in.dims, in.h, in.w, in.c);

    ncnn::Extractor ex = shufflenetv2.create_extractor();

    ex.input("input", in);//input 是 .param文件中输入节点名称

    ncnn::Mat out;
    ex.extract("output", out);

    fprintf(stderr, "output shape: %d %d %d %d\n", out.dims, out.h, out.w, out.c);

    cls_scores.resize(out.w);
    for (int j=0; j<out.w; j++)
    {
        cls_scores[j] = out[j];
        printf("cls_scores[%d]=%f\n",j,cls_scores[j]);
    }

    return 0;
}

static int print_topk(const std::vector<float>& cls_scores, std::vector<int>& index_vec, int topk)
{
    // partial sort topk with index
    int size = cls_scores.size();
    std::vector< std::pair<float, int> > vec;
    vec.resize(size);
    for (int i=0; i<size; i++)
    {
        vec[i] = std::make_pair(cls_scores[i], i);
    }

    std::partial_sort(vec.begin(), vec.begin() + topk, vec.end(),
                      std::greater< std::pair<float, int> >());

    // print topk and score
    for (int i=0; i<topk; i++)
    {
        float score = vec[i].first;
        int index = vec[i].second;
        index_vec.push_back(index);
        fprintf(stderr, "%d = %f\n", index, score);
    }

    return 0;
}



int main(int argc, char** argv)
{


    const char* imagepath = argv[1];

    cv::Mat m = cv::imread("/home/ubuntu/CLionProjects/untitled2/0abad4440c4415b9707f11151762526c.jpg");

    // show results
    std::vector<std::string> labels;
   

    std::vector<float> cls_scores;

    detect_squeezenet(m, cls_scores);



    std::vector<int> index;
    // print_topk(cls_scores, index, 3);
    for (int i = 0; i < index.size(); i++)
    {
        // cv::putText(m, labels[index[i]], cv::Point(50, 50 + 30 * i), CV_FONT_HERSHEY_SIMPLEX, 1.2, cv::Scalar(0, 100, 200), 2, 8);
    }

    //cv::imshow("m", m);
   // cv::imwrite("test_result.jpg", m);
    //wcv::waitKey(0);

    return 0;
}

测试结果

/home/ubuntu/CLionProjects/untitled2/cmake-build-debug/main
input shape: 3 649 1202 3
output shape: 1 1 10 1
cls_scores[0]=-0.937588
cls_scores[1]=-1.634460
cls_scores[2]=0.439976
cls_scores[3]=-0.885857
cls_scores[4]=-1.280403
cls_scores[5]=0.772008
cls_scores[6]=-1.150973
cls_scores[7]=3.494934
cls_scores[8]=2.941223
cls_scores[9]=-1.699838

测试结果

如果不去掉predict.py中的transform中的resize

参考

​​GitHub - Randl/ShuffleNetV2-pytorch: Implementation of ShuffleNetV2 for pytorch​​