El uso del modelo entrenado por pytorch en onnx y mnn (llamada python + cpp)

Pytorch en sí mismo es más problemático de implementar, generalmente usar onnx y mnn es más práctico

El código para entrenar el modelo:

import torch
import torch.nn  as  nn
import torchvision
import torchvision.transforms as transforms
import torch.optim as optim
from torch.optim import lr_scheduler
import torch.onnx


if __name__ == '__main__':

    # transforms
    transform = transforms.Compose([transforms.ToTensor(),
                transforms.Normalize((0.5,), (0.5,))])
   # datasets
    trainset = torchvision.datasets.FashionMNIST('data',download=True,train=True, transform=transform)
    testset = torchvision.datasets.FashionMNIST('data',download=True,train=False,transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=1,
                                            shuffle=True, num_workers=2)
    testloader = torch.utils.data.DataLoader(testset, batch_size=1,
                                            shuffle=False, num_workers=2)

    classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat','Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle Boot')
    print('trainloader--------',trainloader)
    class FashionMNIST(nn.Module):
        def __init__(self,num_classes=10):
            super(FashionMNIST,self).__init__()
            self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1)
            self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)
            self.pool = nn.MaxPool2d(kernel_size=2)
            self.drop = nn.Dropout(0.2)
            self.conv3 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1)
            self.fc1 = nn.Linear(in_features=64* 3 * 3, out_features=128)
            self.fc2 = nn.Linear(in_features=128, out_features=num_classes)
        def forward(self, input):
            output = self.conv1(input)
            output=nn.ReLU()(output)
            output= self.pool(output)
            output = self.conv2(output)
            output=nn.ReLU()(output)
            output= self.pool(output)
            output = self.conv3(output)
            output=nn.ReLU()(output)
            output= self.pool(output)
            output=self.drop(output)
            output = output.view(-1, 64* 3 * 3)
            output=nn.Flatten()(output)
            output = self.fc1(output)
            output = self.fc2(output)
            return  output

    model = FashionMNIST().to(device='cpu')
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    criterion = nn.CrossEntropyLoss()
    exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

    for epoch in range(1):
        for index,data in enumerate(trainloader):
            inputs,labels=data
            #print(labels)
            inputs = inputs.to('cpu')
            labels = labels.to('cpu')
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            if index%1000==0:
                print(index,loss)
    torch.save(model, 'FashionMNIST.pth')
    model.eval()
    x = torch.randn(1, 1, 28, 28, requires_grad=True)
    torch_out = model(x)
    torch.onnx.export(model, x,  "FashionMNIST.onnx", export_params=True, opset_version=10,          
                  do_constant_folding=True,   input_names = ['input'],   output_names = ['output'], 
)

Se guardará un modelo onnx. El modelo pytorch no se puede transferir directamente a mnn, y debe transferirse a onnx antes de

Vale la pena señalar que cuando torch.onnx.export guarda el modelo onnx, es mejor no usar dynamic_axes, porque el valor predeterminado es arreglar la forma y el tamaño de entrada, y es más problemático llamar al modelo mnn más tarde.

Luego convierta el modelo onnx a mnn

mnnconvert -f ONNX  --modelFile FashionMNIST.onnx --MNNModel FashionMNIST.mnn

Llamar a modelos onnx y mnn

import onnxruntime
import MNN
import numpy as np

if __name__ == '__main__':
    x=np.ones([1, 1, 28, 28]).astype(np.float32)
    #onnx
    session = onnxruntime.InferenceSession("FashionMNIST.onnx")
    inputs = {session.get_inputs()[0].name: x}
    outs = session.run(None, inputs)
    print('onnx result is:',outs)
    #mnn
    interpreter = MNN.Interpreter("FashionMNIST.mnn")
    print("mnn load")
    mnn_session = interpreter.createSession()
    input_tensor = interpreter.getSessionInput(mnn_session)
    print(input_tensor.getShape())
    tmp_input = MNN.Tensor((1, 1, 28, 28),\
    MNN.Halide_Type_Float, x[0], MNN.Tensor_DimensionType_Tensorflow)
    print(tmp_input.getShape())
    #print(tmp_input.getData())
    print(input_tensor.copyFrom(tmp_input))
    input_tensor.printTensorData()
   
    interpreter.runSession(mnn_session)
    output_tensor = interpreter.getSessionOutput(mnn_session,'output')
    output_tensor.printTensorData()
    output_data=np.array(output_tensor.getData())
    print('mnn result is:',output_data)
    

resultado de la operación:

C ++ llama a mnn

#include <iostream>
#include<opencv2/core.hpp>
#include<opencv2/imgproc.hpp>
#include<opencv2/highgui.hpp>
#include<MNN/Interpreter.hpp>
#include<MNN/ImageProcess.hpp>

using namespace std;
using namespace cv;
using namespace MNN;
int main()
{   auto net = std::shared_ptr<MNN::Interpreter>(MNN::Interpreter::createFromFile("FashionMNIST.mnn"));//创建解释器
	cout << "Interpreter created" << endl;
	ScheduleConfig config;
	config.numThread = 8;
	config.type = MNN_FORWARD_CPU;
	auto session = net->createSession(config);//创建session
	cout << "session created" << endl;

	auto inTensor = net->getSessionInput(session, NULL);
	auto outTensor = net->getSessionInput(session, NULL);
	auto _Tensor = MNN::Tensor::create<float>(inTensor->shape(), NULL, MNN::Tensor::TENSORFLOW);

	for (int i = 0; i < _Tensor->elementSize(); i++) {
		_Tensor->host<float>()[i] = 1.f;
		
	}
	inTensor->copyFromHostTensor(_Tensor);
	
	//推理
	net->runSession(session);
    auto output= net->getSessionOutput(session, NULL);
	MNN::Tensor feat_tensor(output, output->getDimensionType());
	output->copyToHostTensor(&feat_tensor);
	feat_tensor.print();
	waitKey(0);
	return 0;
}

resultado de la operación:

El resultado es consistente con la llamada de Python

c ++ llamar onnx

#include<onnxruntime_cxx_api.h>
#include<iostream>
#include<vector>
#include<opencv2/core.hpp>
#include<opencv2/imgproc.hpp>
#include<opencv2/highgui.hpp>

using namespace std;
using namespace cv;

static constexpr const int width_ = 28;
static constexpr const int height_ = 28;
static constexpr const int channel = 1;

std::array<float, width_* height_* channel> input_image_{};
std::array<float, 10> results_{};
int result_{ 0 };

Ort::Value input_tensor_{ nullptr };
std::array<int64_t, 4> input_shape_{ 1,1, width_, height_ };

Ort::Value output_tensor_{ nullptr };
std::array<int64_t, 2> output_shape_{ 1, 10};

OrtSession* session_ = nullptr;
OrtSessionOptions* session_option;

int main() {
	Ort::Env env{ ORT_LOGGING_LEVEL_WARNING, "test" };
	auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
	input_tensor_ = Ort::Value::CreateTensor<float>(memory_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size());

	const char* input_names[] = { "input" };
	const char* output_names[] = { "output" };
	Ort::SessionOptions session_option;
	session_option.SetIntraOpNumThreads(1);
	session_option.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);

	Ort::Session session_(env, L"FashionMNIST.onnx", session_option);
	cout << "session created sucess" << endl;

	float* output = input_image_.data();
	fill(input_image_.begin(), input_image_.end(), 1.f);
	
	//执行推理
	session_.Run(Ort::RunOptions{ nullptr }, input_names, &input_tensor_, 1, output_names, &output_tensor_, 1);
	for (int i = 0;i<10;i++) {
		cout << output_tensor_.GetTensorMutableData<float>()[i] << endl;
	}
	return 0;
}

resultado de la operación:

Resultados consistentes