在 C++ 平臺上部署 PyTorch 模型流程+踩坑實錄
共 9162字，需瀏覽 19分鐘
·
2024-05-25 13:07
點擊上方“小白學視覺”，選擇加"星標"或“置頂”
重磅干貨，第一時間送達
來自 | 知乎 作者丨火星少女
鏈接丨h(huán)ttps://zhuanlan.zhihu.com/p/146453159
1.模型轉換
libtorch不依賴于python，python訓練的模型，需要轉換為script model才能由libtorch加載，并進行推理。在這一步官網提供了兩種方法：
方法一：Tracing
這種方法操作比較簡單，只需要給模型一組輸入，走一遍推理網絡，然后由torch.ji.trace記錄一下路徑上的信息并保存即可。示例如下：
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)
方法二：Scripting
直接在Torch腳本中編寫模型并相應地注釋模型，通過torch.jit.script編譯模塊，將其轉換為ScriptModule。示例如下：
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（https://pytorch.org/docs/master/generated/torch.jit.ignore.html#torch.jit.ignore）
或者 @torch.jit.unused（https://pytorch.org/docs/master/generated/torch.jit.unused.html#torch.jit.unused）
# Same behavior as pre-PyTorch 1.2
@torch.jit.script
def some_fn():
return 2
# Marks a function as ignored, if nothing
# ever calls it then this has no effect
@torch.jit.ignore
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.
@torch.jit.unused
def some_fn3():
import pdb; pdb.set_trace()
return 4
# Doesn't do anything, this function is already
# the main entry point
@torch.jit.export
def some_fn4():
return 2
在這一步遇到好多坑，主要原因可歸為一下兩點
1. 不支持的操作
TorchScript支持的操作是python的子集，大部分torch中用到的操作都可以找到對應實現(xiàn)，但也存在一些尷尬的不支持操作，詳細列表可見https://pytorch.org/docs/master/jit_unsupported.html#jit-unsupported，下面列一些我自己遇到的操作：
1）參數(shù)/返回值不支持可變個數(shù)，例如
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
2)各種iteration操作
eg1.
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])
eg2.
seq_iter = enumerate(scores)
try:
_, inivalues = seq_iter.__next__()
except:
_, inivalues = seq_iter.next()
eg3.
line = next(infile)
3)不支持的語句
eg1. 不支持continue
torch.jit.frontend.UnsupportedNodeError: continue statements aren’t supported
eg2. 不支持try-catch
torch.jit.frontend.UnsupportedNodeError: try blocks aren’t supported
eg3. 不支持with語句
4）其他常見op/module
eg1. torch.autograd.Variable
解決：使用torch.ones/torch.randn等初始化+.float()/.long()等指定數(shù)據(jù)類型。
eg2. torch.Tensor/torch.LongTensor etc.
解決：同上
eg3. requires_grad參數(shù)只在torch.tensor中支持，torch.ones/torch.zeros等不可用
eg4. tensor.numpy()
eg5. tensor.bool()
解決：tensor.bool()用tensor>0代替
eg6. self.seg_emb(seg_fea_ids).to(embeds.device)
解決：需要轉gpu的地方顯示調用.cuda()
總之一句話：除了原生python和pytorch以外的庫，比如numpy什么的能不用就不用，盡量用pytorch的各種API。
2. 指定數(shù)據(jù)類型
1）屬性，大部分的成員數(shù)據(jù)類型可以根據(jù)值來推斷，空的列表/字典則需要預先指定
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())
2）常量，使用Final關鍵字
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())
3）變量。默認是tensor類型且不可變，所以非tensor類型必須要指明
def forward(self, batch_size:int, seq_len:int, use_cuda:bool):
方法三：Tracing and Scriptin混合
一種是在trace模型中調用script，適合模型中只有一小部分需要用到控制流的情況，使用實例如下：
import torch
@torch.jit.script
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)))
另一種情況是在script module中用tracing生成子模塊，對于一些存在script module不支持的python feature的layer，就可以把相關layer封裝起來，用trace記錄相關layer流，其他layer不用修改。使用示例如下：
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())
2.保存序列化模型
如果上一步的坑都踩完，那么模型保存就非常簡單了，只需要調用save并傳遞一個文件名即可，需要注意的是如果想要在gpu上訓練模型，在cpu上做inference，一定要在模型save之前轉化，再就是記得調用model.eval(),形如
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")
3.C++ load訓練好的模型
libtorch/
bin/
include/
lib/
share/
然后就可以構建應用程序了，一個簡單的示例目錄結構如下：
example-app/
CMakeLists.txt
example-app.cpp
example-app.cpp和CMakeLists.txt的示例代碼分別如下：
#include <torch/script.h> // One-stop header.
#include <iostream>#include <memory>
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)
至此，就可以運行以下命令從example-app/文件夾中構建應用程序啦：
mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
cmake --build . --config Release
其中/path/to/libtorch是之前下載后的libtorch文件夾所在的路徑。這一步如果順利能夠看到編譯完成100%的提示，下一步運行編譯生成的可執(zhí)行文件，會看到“ok”的輸出，可喜可賀！
4. 執(zhí)行Script Module
終于到最后一步啦！下面只需要按照構建輸入傳給模型，執(zhí)行forward就可以得到輸出啦。一個簡單的示例如下：
// 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';
前兩行創(chuàng)建一個torch::jit::IValue的向量，并添加單個輸入. 使用torch::ones()創(chuàng)建輸入張量，等效于C ++ API中的torch.ones。然后，運行script::Module的forward方法，通過調用toTensor()將返回的IValue值轉換為張量。C++對torch的各種操作還是比較友好的，通過torch::或者后加_的方法都可以找到對應實現(xiàn)，例如
torch::tensor(input_list[j]).to(at::kLong).resize_({batch, 128}).clone()
//torch::tensor對應pytorch的torch.tensor; at::kLong對應torch.int64;resize_對應resize
最后check一下確保c++端的輸出和pytorch是一致的就大功告成啦~
踩了無數(shù)坑，薅掉了無數(shù)頭發(fā)，很多東西也是自己一點點摸索的，如果有錯誤歡迎指正！
參考資料：
PyTorch C++ API - PyTorch master document
Torch Script - PyTorch master documentation
文章地址：
https://pytorch.org/cppdocs/
https://pytorch.org/tutorials/advanced/cpp_export.html
下載1：OpenCV-Contrib擴展模塊中文版教程
在「小白學視覺」公眾號后臺回復：擴展模塊中文教程，即可下載全網第一份OpenCV擴展模塊教程中文版，涵蓋擴展模塊安裝、SFM算法、立體視覺、目標跟蹤、生物視覺、超分辨率處理等二十多章內容。
下載2：Python視覺實戰(zhàn)項目52講 在「小白學視覺」公眾號后臺回復：Python視覺實戰(zhàn)項目，即可下載包括圖像分割、口罩檢測、車道線檢測、車輛計數(shù)、添加眼線、車牌識別、字符識別、情緒檢測、文本內容提取、面部識別等31個視覺實戰(zhàn)項目，助力快速學校計算機視覺。
下載3：OpenCV實戰(zhàn)項目20講 在「小白學視覺」公眾號后臺回復：OpenCV實戰(zhàn)項目20講，即可下載含有20個基于OpenCV實現(xiàn)20個實戰(zhàn)項目，實現(xiàn)OpenCV學習進階。
交流群
歡迎加入公眾號讀者群一起和同行交流，目前有SLAM、三維視覺、傳感器、自動駕駛、計算攝影、檢測、分割、識別、醫(yī)學影像、GAN、算法競賽等微信群（以后會逐漸細分），請掃描下面微信號加群，備注：”昵稱+學校/公司+研究方向“，例如：”張三 + 上海交大 + 視覺SLAM“。請按照格式備注，否則不予通過。添加成功后會根據(jù)研究方向邀請進入相關微信群。請勿在群內發(fā)送廣告，否則會請出群，謝謝理解~