@俠之大者_(dá)7d3f 你好撒踪,下面貼出C++里面的
int main() {
std::shared_ptr<torch::jit::script::Module> ptModule = torch::jit::load("/data/testcode/modelfile.pt");
assert(ptModule != nullptr);
ptModule->to(torch::kCUDA);
cv::Mat srcImage,resImage,traImage,norImage,detImage;
srcImage = cv::imread(strFile, cv::ImreadModes::IMREAD_COLOR);
cv::cvtColor(srcImage, traImage, cv::COLOR_BGR2RGB);
cv::resize(traImage, resImage, cv::Size(224, 224), 0, 0, CV_INTER_AREA);//CV_INTER_LINEAR CV_INTER_AREA
cv::normalize(resImage, norImage, 1, 0, cv::NORM_MINMAX);
cv::convertScaleAbs(norImage, detImage);
at::Tensor tensorImage = torch::from_blob(detImage.data, {1, detImage.rows, detImage.cols, 3}, torch::kByte);
tensorImage = tensorImage.permute({0, 3, 1, 2});
tensorImage = tensorImage.toType(torch::kFloat);
tensorImage = tensorImage.to(torch::kCUDA);
torch::Tensor result = ptModule->forward({tensorImage}).toTensor();
auto max_result = result.max(1,true);
auto max_index = std::get<1>(max_result).item<float>();
printf("detres id=%d\n", max_index);
return 0;
}
下面是模型轉(zhuǎn)換代碼
model = torchvision.models.resnet50(pretrained = True)
in_features= model.fc.in_features
model.fc = nn.Linear(in_features,2)
model.load_state_dict(torch.load('/data/mode/50/vest.pt',map_location='cuda'))
model.eval()
example = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)
output = traced_script_module(torch.ones(1, 3, 224, 224))
traced_script_module.save("/data/testcode/ptc++/modelfile.pt")
pytroch學(xué)習(xí)(二十一)—C++(libTorch)調(diào)用pytroch預(yù)訓(xùn)練模型前言 當(dāng)我們訓(xùn)練好一個(gè)CNN模型之后,可能要集成到項(xiàng)目工程中芬首,或者移植到到不同的開發(fā)平臺(tái)(比如Android, IOS), 一般項(xiàng)目工程或者App大多數(shù)采用C/C++, Ja...