libtorch与pytorch对应关系 libtorch和pytorch的区别

简介

  楼主在这一年从事的图像实时检测中涉及到深度学习的内容，于是在没有积累的情况下自己探索了三种方法，接下来分为三篇文章依次进行介绍并对比分析它们的优点。

  目前使用比较广泛的深度学习框架有pytorch和TensorFlow两种，我选择使用的是pytorch，但原本的工程是基于C++实现的，为了适配于pytorch我发现了libtorch。libtorch可以说是pytorch的C++移植版本，其安装和使用简单，也不需要复杂的环境配置。本文章将介绍如何用libtorch实现构建神经网络并且加载模型进行预测。

libtorch安装

  在pytorch官网我们可以直接下载libtorch的发行版本，这里提供CPU版本的下载地址

https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-1.9.0%2Bcpu.zip

  直接下载压缩包后解压即可

  官方有提供cmake链接libtorch库的教程，大家可以参考教程进行配置，这里也简单介绍以下。

  项目工程的cmakelists下链接libtorch只需要完成下列几个步骤：

1.指定libtorch目录

set(CMAKE_PREFIX_PATH ~/libtorch)

2.搜索libtorch的package

find_package(Torch REQUIRED)

3.链接libtorch库

target_link_libraries(需要链接库的名称 Libtorch)

4.检验是否配置成功

  运行下列cpp代码：

#include <torch/torch.h>
#include <iostream>

int main() {
  torch::Tensor tensor = torch::rand({2, 3});
  std::cout << tensor << std::endl;
}

  若结果为：

0.2063  0.6593  0.0866
0.0796  0.5841  0.1569

  即配置成功。

libtorch构建网络

  这里简要介绍一下libtorch构建神经网络的方法，以简单的LeNet5网络为例：

  调用libtorch我们需要先引入两个头文件

#include <torch/script.h> // One-stop header.
#include <torch/torch.h>

定义网络结构并封装成库

struct LeNet5 : torch::nn::Module
{
    LeNet5(){};
    LeNet5(int num_class, int padding);

    torch::Tensor forward(torch::Tensor X);

    torch::nn::Conv2d C1{nullptr};
    torch::nn::Conv2d C2{nullptr};

    torch::nn::Linear FC1{nullptr};
    torch::nn::Linear FC2{nullptr};
    torch::nn::Linear FC3{nullptr};

    void Train();
    void Test(const string& address);
    void SaveModel(torch::serialize::OutputArchive& archive, const string& path);
};

网络结构的内部实现

/**
 * @brief LeNet5 构造函数
 *        共七层
 * @note  
 */
LeNet5::LeNet5(int num_class, int padding)
{
    this->C1 = register_module("C1", torch::nn::Conv2d(torch::nn::Conv2dOptions(1, 6, {5, 5}).padding(padding)));
    this->C2 = register_module("C2", torch::nn::Conv2d(torch::nn::Conv2dOptions(6, 16, {5, 5})));
    this->FC1 = register_module("FC1", torch::nn::Linear(torch::nn::LinearOptions(400, 120)));
    this->FC2 = register_module("FC2", torch::nn::Linear(torch::nn::LinearOptions(120, 84)));
    this->FC3 = register_module("FC3", torch::nn::Linear(torch::nn::LinearOptions(84, num_class)));
}

/**
 * @brief 前馈函数，计算权重
 */
torch::Tensor LeNet5::forward(torch::Tensor X)
{
    //卷积池化层
    X = this->C1->forward(X);
    X = torch::nn::functional::relu(X);

    //最大池化
    X = torch::max_pool2d(X, {2, 2}, 2);

    X = this->C2->forward(X);
    X = torch::nn::functional::relu(X);

    //最大池化
    X = torch::max_pool2d(X, {2, 2}, 2);

    //全连接层
    X = X.view({X.size(0), -1});
    X = FC1->forward(X);
    X = torch::nn::functional::relu(X);

    X = FC2->forward(X);
    X = torch::nn::functional::relu(X);

    X = FC3->forward(X);
    return X;
}

损失函数和优化器

  损失函数可以在torch::nn中选择，优化器在torch::optim中也有定义，示例：

// 损失函数定义成交叉熵损失
auto criterion = torch::nn::CrossEntropyLoss();
// 优化器选择SGD优化器
auto optimizer = torch::optim::SGD(this->parameters(), torch::optim::SGDOptions(0.005).momentum(0.9));

保存模型

  libtorch的模型需要保存为torch::serialize::OutputArchive类型，所以可以参考一下方法：

// 实例化保存模型类型
torch::serialize::OutputArchive archive;
// 保存模型参数
this->save(archive);
// 设置保存路径
archive.save_to(path);

总结

  libtorch虽然使用简单方便，但是对于深度学习领域C++对比Python本就不占优势，且移植过程也有很多不适配的地方，所以libtorch在一些方面性能上略差于原本的pytorch，对于实时要求很高的检测中并不适用。

参考内容

INSTALLING C++ DISTRIBUTIONS OF PYTORCH