他山之石在C++平台上部署PyTorch模型流程+踩坑实录
Posted 人工智能前沿讲习
tags:
篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了他山之石在C++平台上部署PyTorch模型流程+踩坑实录相关的知识,希望对你有一定的参考价值。
最近因为工作需要,要把pytorch的模型部署到c++平台上,基本过程主要参照官网的教学示例,期间发现了不少坑,特此记录。
01
模型转换
import torch
import torchvision
# An instance of your model.
model = torchvision.models.resnet18()
# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 224, 224)
# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)
class MyModule(torch.nn.Module):
def __init__(self, N, M):
super(MyModule, self).__init__()
self.weight = torch.nn.Parameter(torch.rand(N, M))
def forward(self, input):
if input.sum() > 0:
output = self.weight.mv(input)
else:
output = self.weight + input
return output
my_module = MyModule(10,20)
sm = torch.jit.script(my_module)
-
forward方法会被默认编译,forward中被调用的方法也会按照被调用的顺序被编译 -
如果想要编译一个forward以外且未被forward调用的方法,可以添加 @torch.jit.export. -
如果想要方法不被编译,可使用@torch.jit.ignore[1] 或者 @torch.jit.unused[2]
# Same behavior as pre-PyTorch 1.2
def some_fn():
return 2
# Marks a function as ignored, if nothing
# ever calls it then this has no effect
def some_fn2():
return 2
# As with ignore, if nothing calls it then it has no effect.
# If it is called in script it is replaced with an exception.
def some_fn3():
import pdb; pdb.set_trace()
return 4
# Doesn't do anything, this function is already
# the main entry point
def some_fn4():
return 2
def __init__(self, **kwargs):
或者
if output_flag == 0:
return reshape_logits
else:
loss = self.loss(reshape_logits, term_mask, labels_id)
return reshape_logits, loss
layers = [int(a) for a in layers]
报错torch.jit.frontend.UnsupportedNodeError: ListComp aren’t supported
for k in range(len(layers)):
layers[k] = int(layers[k])
seq_iter = enumerate(scores)
try:
_, inivalues = seq_iter.__next__()
except:
_, inivalues = seq_iter.next()
line = next(infile)
3)不支持的语句
from typing import Dict
class MyModule(torch.nn.Module):
my_dict: Dict[str, int]
def __init__(self):
super(MyModule, self).__init__()
# This type cannot be inferred and must be specified
self.my_dict = {}
# The attribute type here is inferred to be `int`
self.my_int = 20
def forward(self):
pass
m = torch.jit.script(MyModule())
try:
from typing_extensions import Final
except:
# If you don't have `typing_extensions` installed, you can use a
# polyfill from `torch.jit`.
from torch.jit import Final
class MyModule(torch.nn.Module):
my_constant: Final[int]
def __init__(self):
super(MyModule, self).__init__()
self.my_constant = 2
def forward(self):
pass
m = torch.jit.script(MyModule())
def forward(self, batch_size:int, seq_len:int, use_cuda:bool):
方法三:Tracing and Scriptin混合
import torch
def foo(x, y):
if x.max() > y.max():
r = x
else:
r = y
return r
def bar(x, y, z):
return foo(x, y) + z
traced_bar = torch.jit.trace(bar, (torch.rand(3), torch.rand(3), torch.rand(3)))
import torch
import torchvision
class MyScriptModule(torch.nn.Module):
def __init__(self):
super(MyScriptModule, self).__init__()
self.means = torch.nn.Parameter(torch.tensor([103.939, 116.779, 123.68])
.resize_(1, 3, 1, 1))
self.resnet = torch.jit.trace(torchvision.models.resnet18(),
torch.rand(1, 3, 224, 224))
def forward(self, input):
return self.resnet(input - self.means)
my_script_module = torch.jit.script(MyScriptModule())
02
保存序列化模型
gpu_model.eval()
cpu_model = gpu_model.cpu()
sample_input_cpu = sample_input_gpu.cpu()
traced_cpu = torch.jit.trace(traced_cpu, sample_input_cpu)
torch.jit.save(traced_cpu, "cpu.pth")
traced_gpu = torch.jit.trace(traced_gpu, sample_input_gpu)
torch.jit.save(traced_gpu, "gpu.pth")
03
C++ load训练好的模型
libtorch/
bin/
include/
lib/
share/
example-app/
CMakeLists.txt
example-app.cpp
int main(int argc, const char* argv[]) {
if (argc != 2) {
std::cerr << "usage: example-app <path-to-exported-script-module>\n";
return -1;
}
torch::jit::script::Module module;
try {
// Deserialize the ScriptModule from a file using torch::jit::load().
module = torch::jit::load(argv[1]);
}
catch (const c10::Error& e) {
std::cerr << "error loading the model\n";
return -1;
}
std::cout << "ok\n";
}
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)
find_package(Torch REQUIRED)
add_executable(example-app example-app.cpp)
target_link_libraries(example-app "${TORCH_LIBRARIES}")
set_property(TARGET example-app PROPERTY CXX_STANDARD 14)
mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
cmake --build . --config Release
04
执行Script Module
// Create a vector of inputs.
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({1, 3, 224, 224}));
// Execute the model and turn its output into a tensor.
at::Tensor output = module.forward(inputs).toTensor();
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';
torch::tensor(input_list[j]).to(at::kLong).resize_({batch, 128}).clone()
//torch::tensor对应pytorch的torch.tensor; at::kLong对应torch.int64;resize_对应resize
[1] https://pytorch.org/docs/master/generated/torch.jit.ignore.html#torch.jit.ignore
[2] https://pytorch.org/docs/master/generated/torch.jit.unused.html#torch.jit.unused
[3] https://pytorch.org/docs/master/jit_unsupported.html#jit-unsupported
https://pytorch.org/cppdocs/
https://pytorch.org/tutorials/advanced/cpp_export.html
直播预告
历史文章推荐
分享、点赞、在看,给个三连击呗!
以上是关于他山之石在C++平台上部署PyTorch模型流程+踩坑实录的主要内容,如果未能解决你的问题,请参考以下文章
在ubuntu下将Pytorch模型部署到c++(环境搭建)
全流程从头在树莓派4B上部署自己训练的yolov5模型(配合NCS2加速)