由于python的易用性，深度学习模型多是在python框架下进行训练的，如TensorFlow，pytorch等。而由于硬件设备的限制，有时候其部署可能需要基于C++/C的平台。比如在我们的项目中，语义分割网络是在pytorch下训练的，而分割结果基于应用的后处理部分是在C下面实现的，那怎么才能把这两种平台下的东西结合起来一起运行呢？我想到的方法有以下几种：
将pytorch模型训练好后转成caffe的模型，然后利用caffe的接口在C++下面实现模型的推理应用；
将C部分的代码打包编译成一个动态连接库dll，然后在python框架下调用该dll实现c下面的功能；
利用pytorch的C++版本LibTorch实现pytorch模型的调用。
本文主要记录最后一种方法。
LibTorch 的下载及使用
LibTorch 是pytorch的C++版本，在pytorch版本1.0后就有了。在官网通过如下选择，就可以得到下载链接。
下载链接里有release版本和debug版本，建议两个版本都下载，两者主要是对应的dll和lib不一样，debug版本还提供了pdb，可以帮助定位错误位置。将release版本解压后，得到一个LibTorch的文件夹，再将debug版本解压，将其中的lib文件夹改名为lib_debug，同样放在之前release版本解压的LibTorch文件夹，这样就release和debug版本都可以使用了。
下载的LibTorch中提供了cmakelist，在linux平台可以利用cmake来使用它。而如果在Windows平台利用vs，只需要和一般的第三方库使用一样，在对应的工程中添加正确的AdditionalIncludeDirectories，AdditionalLibraryDirectories，AdditionalDependencies等就可以了。我在实验时，将lib里边所有的lib文件都加入到AdditionalDependencies了。程序运行的时候还需要把对应的dll拷贝到exe所在的文件夹。我使用debug时遇到了一个编译错误，添加preprocessor _SCL_SECURE_NO_WARNINGS就好了。
使用流程
利用LibTorch来调用pytorch模型的流程大致是这样的：
- pytorch训练好模型
- 将模型序列化并存成pt文件
- 在C中利用LibTorch的接口进行正向推演
pytorch模型序列化
第一步我们就不介绍了，我们从第二步开始。模型的序列化是利用Torch Script来完成的。TorchScript是一种从PyTorch代码创建可序列化和可优化模型的方法。用TorchScript编写的任何代码都可以从Python进程中保存并加载到没有Python依赖关系的进程中。对于一个已经训练好的pytorch模型，官方提供两种方法进行Torch Script的转换：tracing和annotation。
Tracing
Tracing的方法还是很简单的，参见如下示例代码：
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)
Annotation
tracing适用于大多数网络，如果你的网络的forward方法中对input有逻辑判断，比如input的size为一个值时走向一个分支，而为另一值时走向另一个分支，那么只能用annotation进行转换。比如如下的网络：
import torch
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
利用annotation来将上述网络模型转成Torch Script可以按如下代码：
my_module = MyModule(10,20)
sm = torch.jit.script(my_module)
annotation的方法我并没有测试，我使用的模型用tracing就已经足够了。
序列化
序列化的意思是指将上述Torch Script描述的模型存成一个文件。
traced_script_module.save("traced_resnet_model.pt")
C++中的正向推演
#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;
}
// Simple tests of the model
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';
std::cout << "ok\n";
}
自己训练的模型的实际操作
下面以我们训练的语义分割网络为例，介绍如何将自己的模型在C++中跑起来。
在实际的操作中，也是遇到了一些问题的。
GPU及DataParallel的问题
第一个问题是我们之前的模型训练是在GPU（相信应该都是这样的）中进行的，并且使用了DataParallel，在序列化时，如下代码是正确的，可以与示例代码做下比较。
device = torch.device('cuda')
model = get_model(args_in)
model = torch.nn.DataParallel(model, device_ids=[0])
model.load_state_dict(torch.load(args_in.test_model_path))
model.to(device)
# use evaluation mode to ignore dropout, etc
model.eval()
# The tracing input need not to be the same size as the forward case.
example = torch.rand(1, 3, 1080, 1920).to(device)
# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model.module, example)
traced_script_module.save("traced_model.pt")
对于GPU训练的模型，需要将模型和tracing用的tensor通过to(device)或者.cuda()转到GPU上，如第5,10行。对于利用DataParallel训练的模型，需要在trace时使用model.module，如第13行。
关于DataParallel多说一句，如果希望正向的时候不需要像第3行那样将model再包一层，在训练save model的时候应该按如下
torch.save(model.module.state_dict(), save_path)
这样存的model就不需要第3行代码，而且第13行的.module也不需要了。
附上因为DataParallel没弄对在pycharm中遇到的错误
RuntimeError: hasSpecialCase INTERNAL ASSERT FAILED at ..\torch\csrc\jit\passes\alias_analysis.cpp:300, please report a bug to PyTorch. We don't have an op for aten::to but it isn't a special case. (analyzeImpl at ..\torch\csrc\jit\passes\alias_analysis.cpp:300)
网络输出是Tuple的问题
我们的网络输出是一个tuple而不是一个tensor，于是在C++调用的时候总是crash，用了debug版本的LibTorch，才发现问题。官方提到LibTorch这种方式需要网络的输出是一个tuple或者tensor，那如果输出的是tuple，在C++端代码应该按如下修改
torch::Tensor result = module.forward(input).toTuple()->elements()[0].toTensor();
图像的前处理
在pytorch模型的训练过程中，我们一般会对图像进行一些前处理，比如
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
在LibTorch中，可以这样做
tensor_image = tensor_image.toType(torch::kFloat);
tensor_image = tensor_image.div(255);
// Normalization
tensor_image[0][0] = tensor_image[0][0].sub_(0.485).div_(0.229);
tensor_image[0][1] = tensor_image[0][1].sub_(0.456).div_(0.224);
tensor_image[0][2] = tensor_image[0][2].sub_(0.406).div_(0.225);
最后贴上我们利用opencv读视频，然后对每一帧运行语义分割正向的代码。
// module for forward process
torch::jit::script::Module module;
try {
// Deserialize the ScriptModule from a file using torch::jit::load().
module = torch::jit::load("traced_model.pt");
} catch (const c10::Error &e) {
std::cerr << "error loading the model\n";
}
torch::DeviceType device = torch::kCUDA;
module.to(device);
// opencv windows
cv::namedWindow("Test", 0);
cvMoveWindow("Test", 0, 0);
cv::VideoCapture t_video_in(videoPath);
long nbFrames = static_cast<long>(t_video_in.get(CV_CAP_PROP_FRAME_COUNT));
for (long f = 0; f < nbFrames; f++) {
cv::Mat image, input;
t_video_in >> image;
cv::cvtColor(image, input, CV_BGR2RGB);
// run semantic segmentation to get label image
torch::Tensor tensor_image = torch::from_blob(input.data, { 1, input.rows, input.cols, 3 }, torch::kByte);
tensor_image = tensor_image.permute({ 0, 3, 1, 2 });
tensor_image = tensor_image.toType(torch::kFloat);
tensor_image = tensor_image.div(255);
// Normalization
tensor_image[0][0] = tensor_image[0][0].sub_(0.485).div_(0.229);
tensor_image[0][1] = tensor_image[0][1].sub_(0.456).div_(0.224);
tensor_image[0][2] = tensor_image[0][2].sub_(0.406).div_(0.225);
tensor_image = tensor_image.to(torch::kCUDA);
torch::Tensor result = module.forward({ tensor_image }).toTuple()->elements()[0].toTensor();
torch::Tensor pred = result.argmax(1);
pred = pred.squeeze();
pred = pred.to(torch::kU8);
pred = pred.to(torch::kCPU);
cv::Mat label(cv::Size(image.cols,image.rows), CV_8U, pred.data_ptr());
cv::imshow("Test", label);
cv::waitKey(1);
}
t_video_in.release();
