Python端生成pt模型文件
net. load( model_path)
net. eval ( )
net. to( "cuda" )
example_input = torch. rand( 1 , 3 , 240 , 320 ) . to( "cuda" )
traced_model = torch. jit. trace( net, example_input)
traced_model. save( "model.pt" )
output = traced_model( example_input)
print ( len ( output) )
C++ 端进行调用
c++环境配置 libtorch常用API
# include
# include
# include
int main ( ) {
std:: cout << "cuda::is_available():" << torch:: cuda:: is_available ( ) << std:: endl;
torch:: Tensor tensor = torch:: rand ( { 2 , 3 } ) . to ( at:: kCUDA) ;
std:: cout << tensor << std:: endl;
torch:: jit:: script:: Module module ;
module = torch:: jit:: load ( "/home/yang/Documents/demo/opencv/model.pt" ) ;
std:: vector< torch:: jit:: IValue> inputs;
inputs. push_back ( torch:: rand ( { 1 , 3 , 240 , 320 } ) . to ( at:: kCUDA) ) ;
auto out = module . forward ( inputs) ;
auto tpl = out. toTuple ( ) ;
auto out_ct_hm = tpl-> elements ( ) [ 0 ] . toTensor ( ) ;
out_ct_hm. print ( ) ;
auto out_wh = tpl-> elements ( ) [ 1 ] . toTensor ( ) ;
out_wh. print ( ) ;
}
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38
可能出错的问题
terminate called after throwing an instance of ‘c10::Error’ what(): open file failed, file path: model.pt (FileAdapter at …/…/caffe2/serialize/file_adapter.cc:11)。 模型路径有问题,使用绝对路径解决。 ‘c10::Error’ what(): isTensor() INTERNAL ASSERT FAILED。 很明显,模型的输出应该不是一个 Tensor,可能是一个列表或者元组什么的