如果我想让 OpenCV dnn 模块加载它，我应该如何保存 PyTorch 的模型 [英] How should I save the model of PyTorch if I want it loadable by OpenCV dnn module

问题描述

我通过 PyTorch 训练了一个简单的分类模型并通过 opencv3.3 加载它，但它抛出异常并说

I train a simple classification model by PyTorch and load it by opencv3.3, but it throw exception and say

OpenCV 错误:函数/特性未在 readObject、file 中实现(不支持的 Lua 类型)/home/ramsus/Qt/3rdLibs/opencv/modules/dnn/src/torch/torch_importer.cpp，第 797 行/home/ramsus/Qt/3rdLibs/opencv/modules/dnn/src/torch/torch_importer.cpp:797:错误:(-213) 函数 readObject 中不支持 Lua 类型

模型定义

class conv_block(nn.Module):
    def __init__(self, in_filter, out_filter, kernel):
        super(conv_block, self).__init__()

        self.conv1 = nn.Conv2d(in_filter, out_filter, kernel, 1, (kernel - 1)//2)
        self.batchnorm = nn.BatchNorm2d(out_filter)
        self.maxpool = nn.MaxPool2d(2, 2)

    def forward(self, x):
        x = self.conv1(x)
        x = self.batchnorm(x)
        x = F.relu(x)
        x = self.maxpool(x)

        return x

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        self.conv1 = conv_block(3, 6, 3)
        self.conv2 = conv_block(6, 16, 3)
        self.fc1 = nn.Linear(16 * 8 * 8, 120)
        self.bn1 = nn.BatchNorm1d(120)
        self.fc2 = nn.Linear(120, 84)
        self.bn2 = nn.BatchNorm1d(84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size()[0], -1)
        x = F.relu(self.bn1(self.fc1(x)))
        x = F.relu(self.bn2(self.fc2(x)))
        x = self.fc3(x)
        return x

该模型仅使用 Conv2d、ReLU、BatchNorm2d、MaxPool2d 和 Linear 层，opencv3.3 支持每一层

This model use Conv2d, ReLU, BatchNorm2d, MaxPool2d and Linear layer only, every layers are supported by opencv3.3

我用 state_dict 保存

I save it by state_dict

torch.save(net.state_dict(), 'cifar10_model')

通过 C++ 加载它

std::string const model_file("/home/some_folder/cifar10_model");

std::cout<<"read net from torch"<<std::endl;
dnn::Net net = dnn::readNetFromTorch(model_file);

我想我用错误的方式保存模型，为了使用 OpenCV 加载，保存 PyTorch 模型的正确方法是什么?谢谢

I guess I save the model with the wrong way, what is the proper way to save the model of PyTorch in order to load using OpenCV? Thanks

我用另一种方式保存模型，但也无法加载

I use another way to save the model, but it cannot be loaded either

torch.save(net, 'cifar10_model.net')

这是一个错误吗?还是我做错了什么?

Is this a bug?Or I am doing something wrong?

推荐答案

我找到了答案，opencv3.3 不支持 PyTorch (https://github.com/pytorch/pytorch) 但 pytorch (https:///github.com/hughperkins/pytorch)，真是个大惊喜，我不知道还有其他版本的pytorch存在(看起来像一个死项目，好久没有更新了)，希望他们能提一下他们在 wiki 上支持哪些 pytorch.

I found the answer, opencv3.3 do not support PyTorch (https://github.com/pytorch/pytorch) but pytorch (https://github.com/hughperkins/pytorch), it is a big surprise, I never know there are another version of pytorch exist(looks like a dead project, long time haven't updated), I hope they could mention which pytorch they support on wiki.

这篇关于如果我想让 OpenCV dnn 模块加载它，我应该如何保存 PyTorch 的模型的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！