https://pytorch.apachecn.org/docs/1.2/advanced/cpp_export.html
前期:我们拿到的centernet有带DCN版本的,但是CenterNet源码自带编译的,所以先摒弃它,还没尝试如何转到torch里面,使用dlav032模型先。
1使用下面的方法,注意目前只支持较新版本的pytorch,至少0.4是不可以的;
以下代码我是直接写在/src/lib/detectors/base_detector.py 类初始化里面,
self.model.eval()
# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 512, 512).cuda()
res_model = torchvision.models.resnet18()
# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(self.model,example)
# save model
traced_script_module.save("./model.pt")
另外还需要使得模型是单输出单输出,如果输出矩阵宽高相同forward里面concat操作即可,不同的话,我还没想到什么办法,比如下面这个centernet的修改。
def forward(self, x):
x = self.base(x)
x = self.dla_up(x[self.first_level:])
#x = self.fc(x)
#y = self.softmax(self.up(x))
temp_heads = []
for head in self.heads:
temp_head = self.__getattr__(head)(x)
temp_heads.append(temp_head)
re = torch.cat(temp_heads, 1)
return re
c++工程的验证
int img_size = 512;
std::shared_ptr<torch::jit::script::Module> simaRpnNet;
std::string det_model_path = "/data_1/vir/car_detection/train_ws/CenterNet/c++/model.pt";
torch::NoGradGuard no_grad;
simaRpnNet = torch::jit::load(det_model_path);
simaRpnNet->to(at::kCUDA);
assert(simaRpnNet != nullptr);
cout << "[INFO] init model done...\n";
vector<torch::jit::IValue> inputs; //def an input
cv::Mat src, image, float_image;
src = cv::imread("/data_1/vir/car_detection/train_ws/CenterNet/images/00d13b4a-0306-4d08-a740-b1b5c63f94c40.jpg");
cout<<to_string(src.at<cv::Vec3b>(0, 0)[0])<<endl;
cout<<to_string(src.at<cv::Vec3b>(0, 0)[1])<<endl;
cout<<to_string(src.at<cv::Vec3b>(0, 0)[2])<<endl;
resize(src, image, cv::Size(img_size, img_size), cv::INTER_LINEAR); // resize 图像
cout<<to_string(image.at<cv::Vec3b>(0, 0)[0])<<endl;
cout<<to_string(image.at<cv::Vec3b>(0, 0)[1])<<endl;
cout<<to_string(image.at<cv::Vec3b>(0, 0)[2])<<endl;
// cvtColor(image, image, CV_BGR2RGB); // bgr -> rgb
image.convertTo(float_image, CV_32F, 1.0 / 255); //归一化到[0,1]区间 TODO
float *point_img;
cout<<float_image.at<cv::Vec3f>(0, 0)[0]<<endl;
cout<<float_image.at<cv::Vec3f>(0, 0)[1]<<endl;
cout<<float_image.at<cv::Vec3f>(0, 0)[2]<<endl;
// point_img = float_image.ptr(32);
// cout << *(float_image.data) << endl; //输出一个像素点点值
auto img_tensor = torch::CPU(torch::kFloat32).tensorFromBlob(float_image.data, {1, img_size, img_size, 3}); //将cv::Mat转成tensor,大小为1,224,224,3
img_tensor = img_tensor.permute({0, 3, 1, 2}); //调换顺序变为torch输入的格式 1,3,224,224
//img_tensor[0][0] = img_tensor[0][0].sub_(0.485).div_(0.229); //减去均值,除以标准差
//img_tensor[0][1] = img_tensor[0][1].sub_(0.456).div_(0.224);
//img_tensor[0][2] = img_tensor[0][2].sub_(0.406).div_(0.225);
auto img_var = torch::autograd::make_variable(img_tensor, false); //不需要梯度
inputs.emplace_back(img_var.to(at::kCUDA)); // 把预处理后的图像放入gpu
torch::Tensor result = simaRpnNet->forward(inputs).toTensor(); //前向传播获取结果
inputs.pop_back();
cout<<"Forward over!!!"<<endl;
for (int i = 0; i < 9; ++i ){
cout<<"result:"+ std::to_string(i) +" "<< result[0][i][0][0]<<endl;
}
核对是否正确的时候碰到的几个问题:
1. c++里面输入的时候不做bgr->rgb python工程里面是做了的。
2. 比较后输出矩阵值完全一样