基于C++的PyTorch模型部署

class?MyModel(torch.nn.Module):????def?__init__(self):????????super(MyModel,?self).__init__()        self.linear = torch.nn.Linear(4, 4)
????def?forward(self,?x,?h):????????new_h?=?torch.tanh(self.linear(x)?+?h)        return new_h, new_h
#?創(chuàng)建模型實(shí)例?my_model?=?MyModel()#?輸入示例x,?h?=?torch.rand(3,?4),?torch.rand(3,?4)#?torch.jit.trace方法對(duì)模型構(gòu)建TorchScripttraced_model?=?torch.jit.trace(my_model,?(x,?h))#?保存轉(zhuǎn)換后的模型traced_model.save('model.pt')

print(traced_model.graph)

graph(%self.1?:?__torch__.torch.nn.modules.module.___torch_mangle_1.Module,??????%input?:?Float(3,?4),??????%h?:?Float(3,?4)):??%19?:?__torch__.torch.nn.modules.module.Module?=?prim::GetAttr[name="linear"](%self.1)??%21?:?Tensor?=?prim::CallMethod[name="forward"](%19,?%input)??%12?:?int?=?prim::Constant[value=1]()?#?/var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0??%13?:?Float(3,?4)?=?aten::add(%21,?%h,?%12)?#?/var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0??%14?:?Float(3,?4)?=?aten::tanh(%13)?#?/var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0??%15?:?(Float(3,?4),?Float(3,?4))?=?prim::TupleConstruct(%14,?%14)  return (%15)

print(traced_cell.code)

def?forward(self,????input:?Tensor,????h:?Tensor)?->?Tuple[Tensor,?Tensor]:??_0?=?torch.add((self.linear).forward(input,?),?h,?alpha=1)??_1?=?torch.tanh(_0)  return (_1, _1)

scripted_model?=?torch.jit.script(MyModel)scripted_model.save('model.pt')

import?torchimport?torchvision.models?as?modelsvgg16?=?models.vgg16()example?=?torch.rand(1,?3,?224,?224).cuda()?model?=?model.eval()traced_script_module?=?torch.jit.trace(model,?example)output?=?traced_script_module(torch.ones(1,3,224,224).cuda())traced_script_module.save('vgg16-trace.pt')print(output)

tensor([[?-0.8301,?-35.6095,?12.4716]],?device='cuda:0',        grad_fn=<AddBackward0>)

cmake_minimum_required(VERSION?3.0.0?FATAL_ERROR)project(libtorch_test)find_package(Torch?REQUIRED)message(STATUS?"Pytorch?status:")message(STATUS?"libraries:?${TORCH_LIBRARIES}")add_executable(libtorch_test?test.cpp)target_link_libraries(libtorch_test?"${TORCH_LIBRARIES}")set_property(TARGET libtorch_test PROPERTY CXX_STANDARD 11)

#include?"torch/script.h"#include?"torch/torch.h"#include?#include?using?namespace?std;
int?main(int?argc,?const?char*?argv[]){????if?(argc?!=?2)?{????????std::cerr?<"usage:?example-app?\n";????????return?-1;    }
????//?讀取TorchScript轉(zhuǎn)化后的模型????torch::jit::script::Module?module;????try?{????????module?=?torch::jit::load(argv[1]);    }
????catch?(const?c10::Error&?e)?{????????std::cerr?<"error?loading?the?model\n";????????return?-1;    }
????module->to(at::kCUDA);????assert(module?!=?nullptr);    std::cout << "ok\n";
????//?構(gòu)建示例輸入????std::vector?inputs;    inputs.push_back(torch::ones({1, 3, 224, 224}).to(at::kCUDA));
????//?執(zhí)行模型推理并輸出tensor????at::Tensor?output?=?module->forward(inputs).toTensor();    std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';}

ok-0.8297,?-35.6048,?12.4823[Variable[CUDAFloatType]{1,3}]

import?torchfrom?torchvision.models?import?resnet18model?=resnet18()example?=?torch.rand(1,?3,?224,?224)tracing.traced_script_module = torch.jit.trace(model, example)

traced_script_module.save("traced_resnet_model.pt")

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)

#include??//?torch頭文件.#include #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?{????//?反序列化：導(dǎo)入TorchScript模型????module?=?torch::jit::load(argv[1]);  }
??catch?(const?c10::Error&?e)?{????std::cerr?<"error?loading?the?model\n";????return?-1;  }  std::cout << "ok\n";}

mkdir?example_testcd?example_testcmake?-DCMAKE_PREFIX_PATH=/path/to/libtorch?..cmake --example_test . --config Release

std::vector?inputs;inputs.push_back(torch::ones({1,?3,?224,?224}));//?執(zhí)行推理并將模型轉(zhuǎn)化為Tensoroutput = module.forward(inputs).toTensor();std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';