C++libtorch部署Superpoint+Superglue模型

前言

在项目中需要将pytorch模型部署到C++上使用，得到封装好的Superpoint+Supergluedll与lib文件后，可在其他设备中无需下载环境直接调用相关模型进行推理。
前期尝试了直接将pytorch模型打包为.pt模型或onnx模型的方法，发现若模型中存在自定义模块如Superglue模型中的图神经网络层GNN，则导出的模型无法直接在C++中进行使用（报错程序终止）。踩了很多坑后，在网上发现了一篇关于使用libtorch编写resnet完整模型的文章，借鉴了相关思路后，决定将superglue模型完整用C++语言进行编写。
当然，若后续找到了注册自定义模块到libtorch中地方法，使得导出的模型能够正常地使用，则还是会采用直接调用模型的方法（毕竟不用全部重新写，也便于模型改进）。

2023.7.9 整套系统编写完成，在模型的基础上加入了评估部分与卫星影像检索部分，并进行程序结构的优化，发现release版本相对于debug版本速度提升了3倍。

2023.7.16 进一步优化算法结构，将函数的输入输出都按照opencv的标准进行同一（cv::InputArray、cv::OutputArrary),并把superpoint也重新编写为C++模型，避免实用时用户的CUDA版本不一致导致模型无法使用，本算法初期版本完成。
以下为假期计划推进的部分：

• [1 ] 阅读文献找到优化算法的思路，特别是在高分辨率图像上，原始算法精度较差（算法在低分辨率尺寸下训练，且算法深度无法适应高分辨率图像（8000×6000、4000×3000），导致在resize输入图像的过程中，造成不可不免的像素损失，导致最终定位精度降低）
• [ 2] 进一步学习Libtroch部署，并学习ONNX、Tensorrt部署（感觉自己偏工程，不是搞科研的料）

libtorch环境配置

环境配置无需多言，下好libtorch后，将其中的lib、include、bin文件夹链接到项目中就行。下列是我遇到的相关问题及解决方法：

1. 使用GPU进行推理
   完成环境配置后，会发现torch::cuda::is_available()还是为false，将属性-链接器-命令行中加入下列内容后，即可正确调用GPU

/INCLUDE:?warp_size@cuda@at@@YAHXZ  /INCLUDE:?warp_size@cuda@at@@YAHXZ 
/INCLUDE:?_torch_cuda_cu_linker_symbol_op_cuda@native@at@@YA?
AVTensor@2@AEBV32@@Z 

Superpoint部署

模型导出

关于导出模型的资料网上很多，不再赘述。

device ='cuda'
model = SuperPoint(config.get('superpoint', {})).to(device)
data = torch.randn(1, 1, 480, 640).to(device)
traced_model = torch.jit.trace(model.eval(), data)
torch.jit.save(traced_model, 'superpoint.pt')

以下记录注意事项：

1. 模型和模型的输入设备需相同，且在C++中调用时，使用的设备需和模型一致（GPU或CPU），否则会出现c10::error等错误
2. libtorch无法直接识别自定义模块或函数，如下列非极大值抑制函数：

def simple_nms(scores, nms_radius: int):
    """ Fast Non-maximum suppression to remove nearby points """
    assert (nms_radius >= 0)

    def max_pool(x):
        return torch.nn.functional.max_pool2d(
            x, kernel_size=nms_radius * 2 + 1, stride=1, padding=nms_radius)

    zeros = torch.zeros_like(scores)
    max_mask = scores == max_pool(scores)
    for _ in range(2):
        supp_mask = max_pool(max_mask.float()) > 0

        supp_scores = torch.where(supp_mask, zeros, scores)
        new_max_mask = supp_scores == max_pool(supp_scores)
        max_mask = max_mask | (new_max_mask & (~supp_mask))
    return torch.where(max_mask, scores, zeros)

故导出的模型不能包含类似的自定义函数或模块，只能将模型部分导出，使用到的自定义函数用C++语言重新编写，若后续找到能正确识别的方法，会在此处更新
3.虽然定义模型输入的代码data = torch.randn(1, 1, 480, 640).to(device)形状是固定的，但实际上使用起来是可以变化的

C++导入模型进行推理，并编写相关自定义函数

导入部分代码如下：

	// load model
	auto module = torch::jit::load("./superpoint_gpu.pt");
	module.to(torch::kCUDA);
	module.eval();
	// inference
	auto outputs = module.forward({ image });
	auto tu_outputs = outputs.toTuple();
	sp.scores = tu_outputs->elements()[0].toTensor();
	sp.descriptors = tu_outputs->elements()[1].toTensor();

相关自定义函数就按照python代码中的功能复现就行，此处不作展示。
以下记录注意事项：

1. 模型的输出为多个时，需要先转换为元组toTuple()，再根据索引转换为张量toTensor()；若输出只有一个，auto outputs = module.forward({ image }).toTensor()即可
2. 模型的设备（torch::kCUDA或torch::kCPU）需要和导出模型时一致，否则会报错c10::error

Superglue部署

将各自定义模块在C++中复现

Superpoint模型由于结构较为简单，故部署起来很快，但Superglue模型由于都是自定义的模块，libtorch无法直接识别这些模块，故此处用C++语言重新搭建模型结构。以下为模型中特征编码模块的复现：

//特征点编码器模块
class KeypointEncoderImpl : public torch::nn::Module 
{
    public:
        KeypointEncoderImpl(int feature_dim, std::vector<int64_t> layers)
        { 
        std::vector<int64_t> channels = { 3 };
        channels.insert(channels.end(), layers.begin(), layers.end());
        channels.push_back(feature_dim);
      
        this->encoder = register_module("encoder", MLP(channels));
        torch::nn::init::constant_(this->encoder->named_parameters()["12.bias"], 0.0);
        }
        torch::Tensor forward(torch::Tensor kpts, torch::Tensor scores)
        {
            kpts = kpts.transpose(1, 2);
            scores = scores.unsqueeze(1);
        
            torch::Tensor inputs = my_cat({ kpts, scores }, 1);
            return this->encoder->forward(inputs);
        }
    private: torch::nn::Sequential encoder;
};

TORCH_MODULE(KeypointEncoder);

以下为Superglue模型的部分代码：

class SuperGlueImpl: public torch::nn::Module
{
public:
    SuperGlueImpl()
    {
        std::vector<int64_t> keypoint_encoder = { 32, 64, 128, 256 };
        std::vector<std::string> GNN_layers(18, "");
        for (int i = 0; i < 18; i += 2) {
            GNN_layers[i] = "self";
            GNN_layers[i + 1] = "cross";
        }
        this->kenc = KeypointEncoder(256, keypoint_encoder);
        this->gnn = AttentionalGNN(256, GNN_layers);
        this->final_proj = torch::nn::Conv1d(torch::nn::Conv1dOptions(256, 256, 1).bias(true));
        /*auto bin_score = torch::tensor(4.4124);*/
        bin_score = torch::ones({});
        register_parameter("bin_score", bin_score);
        register_module("kenc", kenc);
        register_module("gnn", gnn);
        register_module("final_proj", final_proj);
    }
private:
    torch::Tensor bin_score;
    KeypointEncoder kenc = nullptr;
    AttentionalGNN gnn = nullptr;
    torch::nn::Conv1d final_proj = nullptr;
};

以下记录注意事项：

1. 在各模块中，需要将子模块注册到当前模型中，便于后续加载权重时，根据子模块参数名进行索引加载权重。register_module是libTorch提供的一个方法，它允许将模块注册为模型的一部分，以便在训练过程中自动跟踪参数和梯度更新。第一个参数是子模块的名称，第二个参数是要注册的实际模块。上述代码中，将encoder注册为KeypointEncoder模块的子模块（由MLP层组成），便于在该模块的前向传播中使用。
2. 完成各模块的编写后，使用TORCH_MODULE(KeypointEncoder);将各模块注册到libtorch中，便于后续Superglue模型调用各模块。
3. 在Superglue模型中，将需要使用的模块使用register_module()定义，注意参数名要和原模型权重保持一致。单独的需要加载权重的参数如bin_score，则使用 register_parameter()将其登记为模型参数。

载入、保存权重

原模型的权重字典如下：

先将官方的.pth权重模型转换为.bin二进制格式，再在C++中根据参数名读取该二进制文件，完成权重的加载。
加载完权重后，将权重保存为C++模型能正确载入的权重文件，后续可直接读取

torch::save(model, "libtorch_superglue.pt");//保存权重
torch::load(model, "libtorch_superglue.pt");//加载权重

推理

  	SuperGlue model;//定义模型
    //加载模型参数与缓存权重
    torch::load(model, "libtorch_superglue.pt");
    model->to(torch::kCUDA);//使用GPU
    model->eval();//评估模式
    auto output = model->forward() //前向传播

整套系统效果如下：

评估部分

根据无人机位姿信息检索大致卫星区域

根据无人机站点经纬度信息，使用gdal库在大范围卫星影像上检索相应的区域，作为算法的输入

根据目标点像素坐标得到相应的经纬度坐标，完成精度评估

存在问题

1. libtorch中许多函数使用时都会报错，如torch::where、torch::cat、torch::index_select等，我直接用自定义函数复现了上述函数的功能，结果能够顺利运行程序，其中的原因还需要进一步探索。
2. 更好的解决方案是将模型中的自定义模块都注册到libtorch中，使其能够被正确识别，但苦于网上资料很少，需要进一步地学习