opencv-python加载pytorch训练好的onnx格式线性回归模型

news2024/11/22 19:19:28

opencv是一个开源的图形库，有针对java,c++,python的库依赖，它本身对模型训练支持的不好，但是可以加载其他框架训练的模型来进行预测。

这里举一个最简单的线性回归的例子，使用深度学习框架pytorch训练模型，最后保存模型为onnx格式。最后使用opencv-python库来进行加载预测。

这里准备的线性回归模型数据如下：

x_data	y_data
1	2
2	4
3	6
4	?
8	?
10	?
15	?

直观的看，这里其实就是一个 y = 2x的线性方程。但是机器学习里面，它是通过不断迭代的方式求得最终的系数，w,b的值。

本文使用过的python库版本:

torch:1.13.0
opencv-python: 4.5.5
numpy 1.24.2
python: 3.10

show me the code:

import torch

x_data = torch.Tensor([[1.0], [2.0], [3.0]])
y_data = torch.Tensor([[2.0], [4.0], [6.0]])


class LinearModel(torch.nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()
        self.linear = torch.nn.Linear(1, 1)

    def forward(self, x):
        return self.linear(x)


model = LinearModel()
criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=0.02)

for epoch in range(1000):
    y_pred = model(x_data)
    loss = criterion(y_pred, y_data)
    if epoch % 100 == 0:
        print(epoch + 1, loss)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

print('w = {}'.format(model.linear.weight.item()))
print('b = {}'.format(model.linear.bias.item()))

x_test = torch.Tensor([[4.0]])
y_test = model(x_test)
print('y_pred(4.0) = ', y_test.data)

x_test = torch.Tensor([[8.0]])
y_test = model(x_test)
print('y_pred(8.0) = ', y_test.data)

x_test = torch.Tensor([[10.0]])
y_test = model(x_test)
print('y_pred(10.0) = ', y_test.data)

x_test = torch.Tensor([[15.0]])
y_test = model(x_test)
print('y_pred(15.0) = ', y_test.data)

model.eval()
dummy_input = torch.randn(1, 1)
input_name = ["input"]
output_name = ["output"]
onnx_name = "test.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_name,
    verbose=True,
    input_names=input_name,
    output_names=output_name
)

运行，打印信息如下：

这里重点关注线性回归系数，w = 1.99，b = 1.31。这里的w其实很接近2，毕竟是计算机算出来的，这里没有明确表示使用梯度下降算法，但是实际上在训练那部分，确实使用的是梯度下降法。

for epoch in range(1000):
    y_pred = model(x_data)
    loss = criterion(y_pred, y_data)
    if epoch % 100 == 0:
        print(epoch + 1, loss)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

最后，通过这个回归系数预测，符合预期。

算法最终会将模型保存在本地项目路径下的test.onnx文件中。

上面算是模型准备好了，下面我们通过opencv-python来加载这个模型，并预测：

from cv2 import dnn
import numpy as np

net = dnn.readNetFromONNX("test.onnx")
matblob = np.full((1, 1), 1024, dtype=np.int32)
net.setInput(matblob)
print('input = {}'.format(matblob))
output = net.forward()
print('output = {}'.format(output))

这里，matblob其实就是一个Mat，但是在opencv-python里面，它可以通过numpy来创建，这里在网上的都是通过读取一个图片来生成matblob对象，我个人觉着我们这里很明确，就是需要指定一个数字1024，我们通过np.full() 来设置cols,rows,type,value就成功创建了这个输入Mat对象。

运行这个代码，我们期望得到的是2048，实际结果如下所示：