开发笔记:他山之石在C++平台上部署PyTorch模型流程+踩坑实录

import torchimport 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)

# Same behavior as pre-PyTorch 1.2@torch.jit.scriptdef some_fn(): return 2
# Marks a function as ignored, if nothing# ever calls it then this has no effect@torch.jit.ignoredef 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.unuseddef some_fn3(): import pdb; pdb.set_trace() return 4
# Doesn't do anything, this function is already# the main entry point@torch.jit.exportdef some_fn4(): return 2

def __init__(self, **kwargs):

if output_flag == 0: return reshape_logitselse: loss = self.loss(reshape_logits, term_mask, labels_id) return reshape_logits, loss

layers = [int(a) for a in layers]

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)

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 Finalexcept: # 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):

import torch
@torch.jit.scriptdef 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 torchimport 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())

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")

libtorch/ bin/ include/ lib/ share/

example-app/ CMakeLists.txt example-app.cpp

#include // One-stop header.#include <iostream>#include int main(int argc, const char* argv[]) { if (argc != 2) { std::cerr <<"usage: example-app \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 buildcd buildcmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..cmake --build . --config Release

// Create a vector of inputs.std::vector 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 </*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