lenet在pytorch迁移

1.mobilenet_v2 模型结构

多话不说先上代码

import torch.nn as nn
import math


def conv_bn(inp, oup, stride):
    return nn.Sequential(
        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
        nn.BatchNorm2d(oup),
        nn.ReLU6(inplace=True)
    )


def conv_1x1_bn(inp, oup):
    return nn.Sequential(
        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
        nn.BatchNorm2d(oup),
        nn.ReLU6(inplace=True)
    )


class InvertedResidual(nn.Module):
    def __init__(self, inp, oup, stride, expand_ratio):
        super(InvertedResidual, self).__init__()
        self.stride = stride
        assert stride in [1, 2]

        hidden_dim = round(inp * expand_ratio)
        self.use_res_connect = self.stride == 1 and inp == oup

        if expand_ratio == 1:
            self.conv = nn.Sequential(
                # dw
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )
        else:
            self.conv = nn.Sequential(
                # pw
                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True),
                # dw
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)


class MobileNetV2(nn.Module):
    def __init__(self, n_class=1000, input_size=224, width_mult=1.):
        super(MobileNetV2, self).__init__()
        block = InvertedResidual
        input_channel = 32
        last_channel = 1280
        interverted_residual_setting = [
            # t, c, n, s
            [1, 16, 1, 1],
            [6, 24, 2, 2],
            [6, 32, 3, 2],
            [6, 64, 4, 2],
            [6, 96, 3, 1],
            [6, 160, 3, 2],
            [6, 320, 1, 1],
        ]

        # building first layer
        assert input_size % 32 == 0
        input_channel = int(input_channel * width_mult)
        self.last_channel = int(last_channel * width_mult) if width_mult > 1.0 else last_channel
        self.features = [conv_bn(3, input_channel, 2)]
        # building inverted residual blocks
        for t, c, n, s in interverted_residual_setting:
            output_channel = int(c * width_mult)
            for i in range(n):
                if i == 0:
                    self.features.append(block(input_channel, output_channel, s, expand_ratio=t))
                else:
                    self.features.append(block(input_channel, output_channel, 1, expand_ratio=t))
                input_channel = output_channel
        # building last several layers
        self.features.append(conv_1x1_bn(input_channel, self.last_channel))
        # make it nn.Sequential
        self.features = nn.Sequential(*self.features)
        self.avgpool = nn.AvgPool2d(kernel_size=4,stride=1)# 7 for 224, 4 for 128
        # building classifier
        self.classifier1 = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(self.last_channel, n_class),
        )

        self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        # print("the x shape :{}".format(x.shape))#(1,1280,4,4)
        # x = x.mean(3).mean(2)#在第3,2个维度求平均值，ncnn不支持mean ,可换成avgpool
        x = self.avgpool(x)
        # print("the x shape :{}".format(x.shape))#(1,1280)
        x = self.classifier1(x.view(-1,1280))
        # print("the x shape :{}".format(x.shape))#(1,2)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                n = m.weight.size(1)
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

利用上面的模型结构训练好一个模型(工作要求我训练了一个人脸正脸与侧脸的检测模型)
链接: [https://pan.baidu.com/s/1W4rZCYzkTstqjKY79o4MHw]
提取码: qb7g

2.加载训练好的mobilenet_v2模型代码

该脚本作用：

1.加载训练好的模型

2.利用pytorch进行预测，预测的结果为

3.到处onnx模型（得到mobilenet_v2_2.onnx模型文件）

链接: [https://pan.baidu.com/s/1S8Vq-5aGVJxzzQHjI3w6Fg]

提取码: 8vrr

4.利用onnx模型进行预测得到的结果为

import torch
# import trochvision
import torch.utils.data
import argparse
import onnxruntime
from mobilenet_v2 import MobileNetV2
import os
import cv2
import numpy as np
from torch.autograd import Variable
from onnxruntime.datasets import get_example


def main(args):
    # print("the version of torch is {}".format(torch.__version__))
    dummy_input=getInput(args.img_size)#获得网络的输入
    # 加载模型
    model = MobileNetV2(n_class=2)
    model_dict =  model.state_dict()
    if args.model_path:
        if os.path.isfile(args.model_path):
            print(("=> start loading checkpoint '{}'".format(args.model_path)))
            state_dict = torch.load(args.model_path)
            print("the best acc is {} in epoch:{}".format(
                state_dict['epoch_acc'], state_dict['epoch']))
            params = state_dict["model_state_dict"]
            # params={k:v for k,v in state_dict.items() if k in  model_dict.keys()}
            # model_dict.update(params)
            # model.load_state_dict(model_dict)
            model.load_state_dict(params)
            print("load cls model successfully")
        else:
            print(("=> no checkpoint found at '{}'".format(args.model_path)))
            return
    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('./test.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,"./mobilenet_v2_2.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,"./mobilenet_v2_2.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="./weights_60/best_epoch_10.tar",
                        help="For training from one model_file")
    parser.add_argument('--save_model_path', type=str, default="./weights_60/best_epoch_10.tar",
                        help="For training from one model_file")
    parser.add_argument('--onnx_model_path', type=str, default="./weights_60/best_epoch_10.tar",
                        help="For training from one model_file")
    parser.add_argument('--img_size', type=int, default=128,
                        help="the image size of model input")
    args = parser.parse_args()
    main(args)

3.下载ncnn代码并编译

3.1 得到mobilenet_v2模型后，需要下载ncnn项目

git clone https://github.com/Tencent/ncnn.git

cd ncnn

mkdir -p build

cd build

cmake …

make -j8

3.2 得到onnx2ncnn运行程序

经过上述一系列的操作后可以在ncnn/tools/onnx/路径下得到onnx2ncnn.cpp文件

同时在ncnn/build/tools/onnx/目录下得到运行文件onnx2ncnn,如下图，运行这个onnx2ncnn文件就可以将onnx模型转化为ncnn模型了（但是先不要着急转化，由于得到的onnx转换模型时有一些冗余，我们用onnx-simplifier 工具简化onnx模型）

4.将onnx转化为ncnn模型

4.1 安装 onnx-simplifier

pip3 install onnx-simplifier

4.2 利用onnx-simplifier简化onnx模型

在终端运行一下命令能够得到简化后的mobilenet_v2_2-sim.onnx文件

python3 -m onnxsim mobilenet_v2_2.onnx mobilenet_v2_2-sim.onnx

4.3 mobilenet_v2_2-sim.onnx文件转ncnn文件

./onnx2ncnn mobilenet_v2_2-sim.onnx mobilenet_v2_2.param mobilenet_v2_2.bin

在终端运行以上命令，可以在ncnn/build/tools/onnx/目录下得到mobilenet_v2_2.param文件和mobilenet_v2_2.bin文件，上图可以看得到，这样就得到了我们梦寐以求的mobilenet_v2的ncnn模型文件了

5.对比ncnn的与onnx与pytorch的预测结果

5.1 利用ncnn模型进行预测

5.1 在ncnn/examples 目录下新建一个mobilenet_v2.cpp文件，然后参考squeezenet的例子，简单修改一下

//
// Created by zjw on 19-6-28.
//

#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 squeezenet;

#if NCNN_VULKAN
    squeezenet.opt.use_vulkan_compute = true;
#endif // NCNN_VULKAN

    squeezenet.load_param("mobilenet_v2_2.param");
    squeezenet.load_model("mobilenet_v2_2.bin");

    ncnn::Mat in = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, bgr.cols, bgr.rows, 128, 128);
    //按照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 = squeezenet.create_extractor();
    
    ex.input("input", in);//input 是 .param文件中输入节点名称
    printf("111111111111111111111\n");

    ncnn::Mat out;
    ex.extract("output", out);
    printf("2222222222222222\n");
    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;
}

static int load_labels(std::string path, std::vector<std::string>& labels)
{
    FILE* fp = fopen(path.c_str(), "r");

    while (!feof(fp))
    {
        char str[1024];
        fgets(str, 1024, fp);  //¶ÁÈ¡Ò»ÐÐ
        std::string str_s(str);

        if (str_s.length() > 0)
        {
            for (int i = 0; i < str_s.length(); i++)
            {
                if (str_s[i] == ' ')
                {
                    std::string strr = str_s.substr(i, str_s.length() - i - 1);
                    labels.push_back(strr);
                    i = str_s.length();
                }
            }
        }
    }
    return 0;
}


int main(int argc, char** argv)
{
    if (argc != 2)
    {
        fprintf(stderr, "Usage: %s [imagepath]\n", argv[0]);
        return -1;
    }
    
    const char* imagepath = argv[1];

    cv::Mat m = cv::imread(imagepath, 1);
    if (m.empty())
    {
        fprintf(stderr, "cv::imread %s failed\n", imagepath);
        return -1;
    }

    // show results
    std::vector<std::string> labels;
    // load_labels("synset_words.txt", labels);

#if NCNN_VULKAN
    ncnn::create_gpu_instance();
#endif // NCNN_VULKAN

    std::vector<float> cls_scores;
    printf("$$$$$$$$$$$$$$$$$ %s\n",argv[1]);
    detect_squeezenet(m, cls_scores);

#if NCNN_VULKAN
    ncnn::destroy_gpu_instance();
#endif // NCNN_VULKAN

    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);
    cv::waitKey(0);

    return 0;
}

5.2 修改examples下的CMakeLists.txt， 添加一下我们的例子

add_executable(mobilenetv3 mobilenetv3.cpp)
target_link_libraries(mobilenetv3 ${NCNN_EXAMPLE_LINK_LIBRARIES})

5.3 在ncnn根目录下放开编译examples

add_subdirectory(examples)

5.4 运行上述命令后，在build/examples/ 下就有mobilenet_v2这个例子

5.5 把转换得到的ncnn模型mobilenet_v2_2.param、mobilenet_v2_2.bin和测试图片全都移到该目录下

5.6 再在目录下运行 mobilenet_v2

./mobilenet_v2 /home/hy/ncnn/build/examples/test.jpg

6.结果比较

pytorch模型预测结果为：

onnx的预测结果为：

ncnn预测的结果为：

结果上看精度还是没什么损失的，接下来再多搞一些例子进行测试再来更新一下