首先训练一个简单的图像分类器。代码如下：
import torch.optim as optim
import torch.nn.functional as F
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.onnx
import torchvision
import torchvision.transforms as transforms
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
shuffle=True, num_workers=0)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
shuffle=False, num_workers=0)
classes = ('plane', 'car', 'bird', 'cat',
'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# functions to show an image
def imshow(img):
img = img / 2 + 0.5 # unnormalize
npimg = img.numpy()
plt.imshow(np.transpose(npimg, (1, 2, 0)))
plt.show()
# get some random training images
dataiter = iter(trainloader)
images, labels = dataiter.next()
print(images.shape)
# show images
imshow(torchvision.utils.make_grid(images))
# print labels
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 3)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 12, 3)
self.conv3 = nn.Conv2d(12, 32, 3)
self.fc1 = nn.Linear(32 * 4 * 4, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = F.relu(self.conv3(x))
x = x.view(-1, 32 * 4 * 4)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
net = Net()
net.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
for epoch in range(10): # loop over the dataset multiple times
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data[0].to(device), data[1].to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 2000 == 1999: # print every 2000 mini-batches
print(outputs)
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 2000))
running_loss = 0.0
print('Finished Training')
# 导出网络到ONNX
dummy_input = torch.randn(1, 3, 32, 32).to(device)
torch.onnx.export(net, dummy_input, "torch.onnx")
# 保存网络位TORCHSCRIPT
dummy_input = torch.randn(1, 3, 32, 32).to(device)
traced_cell = torch.jit.trace(net, dummy_input)
traced_cell.save("tests.pth") 根据opencv官方文档中的说明，可以支持以下框架：Caffe，Darknet，Onnx，Tensorflow，Torch等。但是很可惜，没有我用的pytoch，但是根据第三个参考链接中的方法，可以利用ONNX实现曲线救国。首先利用保存模型方法3所示的办法，将网络和参数保存为对应的格式。然后使用opencv提供的Net cv::dnn::readNetFromONNX ( const String & onnxFile )函数读取保存好的网络。代码实现如下：
//测试opencv加载pytorch模型
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
using namespace cv;
using namespace cv::dnn;
#include <fstream>
#include <iostream>
#include <cstdlib>
using namespace std;
int main()
{
String modelFile = "./torch.onnx";
String imageFile = "./dog.jpg";
dnn::Net net = cv::dnn::readNetFromONNX(modelFile); //读取网络和参数
Mat image = imread(imageFile); // 读取测试图片
cv::cvtColor(image, image, cv::COLOR_BGR2RGB);
Mat inputBolb = blobFromImage(image, 0.00390625f, Size(32, 32), Scalar(), false, false); //将图像转化为正确输入格式
net.setInput(inputBolb); //输入图像
Mat result = net.forward(); //前向计算
cout << result << endl;
}
上述代码就是对第一个参考链接的代码进行了简化，且将输入网络的模型从torch改成ONNX格式。
运行结果如下：
[-0.19793352, -4.0697966, 1.2769811, 2.7011304, 0.22390884, 1.9039617, -0.47333384, -0.15912014, 0.32441139, -2.4327304]
利用pytorch官方提供的LibTorch加载训练好的模型和网络
参考链接：
windows+VS2019+PyTorchLib配置使用攻略
C++调用pytorch，LibTorch在win10下的vs配置和cmake的配置
在C ++中加载TORCHSCRIPT模型官网链接
此处首先说明一下将pytroch保存为TORCHSCRIPT的方法有两种，一种是追踪式，另一种是脚本式。具体介绍见官方文档，理论上此方法两种保存方式都行，这里我们使用追踪式的方法。
首先按照第一个参考链接中的方法配置LibTorch环境，在此基础上参考链接2，还有两个地方需要修改：
- 属性->C/C++ ->常规->SDL检查->否。
- 属性->C/C++ ->语言->符号模式->否。
然后复制粘贴其中的示例代码，进行测试，但是我个人在运行的时候ToTensor(image).to(at::kCUDA);这个语句报错了，提示未定义ToTensor()，这句话的功能也很简单，就是将普通图像格式转化为模型输入需要的格式，于是我又根据第二个参考链接将转化代码进行了修改，代码如下：
#include <torch/script.h>
#include <iostream>
#include <opencv2/opencv.hpp>
#include <torch/torch.h>
// 有人说调用的顺序有关系，我这好像没啥用~~
int main()
{
torch::DeviceType device_type;
if (torch::cuda::is_available()) {
std::cout << "CUDA available! Predicting on GPU." << std::endl;
device_type = torch::kCUDA;
}
else {
std::cout << "Predicting on CPU." << std::endl;
device_type = torch::kCPU;
}
torch::Device device(device_type);
//Init model
std::string model_pb = "tests.pth";
auto module = torch::jit::load(model_pb);
module.to(at::kCUDA);
auto image = cv::imread("dog.jpg", cv::ImreadModes::IMREAD_COLOR);
cv::Mat image_transfomed;
cv::resize(image, image_transfomed, cv::Size(32, 32));
// convert to tensort
torch::Tensor tensor_image = torch::from_blob(image_transfomed.data,
{ image_transfomed.rows, image_transfomed.cols,3 }, torch::kByte);
tensor_image = tensor_image.permute({ 2,0,1 });
tensor_image = tensor_image.toType(torch::kFloat);
tensor_image = tensor_image.div(255);
tensor_image = tensor_image.unsqueeze(0);
tensor_image = tensor_image.to(at::kCUDA);
torch::Tensor output = module.forward({ tensor_image }).toTensor();
auto max_result = output.max(1, true);
auto max_index = std::get<1>(max_result).item<float>();
std::cout << output << std::endl;
//return max_index;
return 0;
}
运行结果如下：
CUDA available! Predicting on GPU.
1.0824 -4.6106 1.0189 2.9937 1.4570 1.4964 -1.3164 -0.7753 0.4567 -3.2543
[ CUDAFloatType{1,10} ]
原著链接：https://blog.csdn.net/cai1493105270/article/details/108127290