VTM工程配置libtorch环境

VTM配置libtorch

首先下载libtorch库,可以百度搜索下载地址.下载完成之后将其解压到文件夹中，如下图所示，

 其中include文件夹下是libtorch库的头文件；
lib文件夹下是libtorch库的静态库（lib）和动态链接库（dll）；
share文件夹下有之后cmake配置时要find的.cmake文件，用于在cmake时让项目链接到libtorch库；

之后配置环境变量，如下图所示，将如下几个目录加入到环境变量中

 配置完成之后，就可以在VTM工程中的CMakeLists.txt中添加如下语句，用于cmake时链接到libtorch库

find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")

 最后使用cmake进行编译即可。

问题：

在编译完成后，打开VTM工程，运行时会出现重定义的错误，这是由于在libtorch库和VTM中存在几个相同名称的变量，可以将VTM中同名称的变量重命名，或者可以将libtorch中同名的变量注释掉（亲测可用）。

libtorch的简单使用：

1. 加载模型

定义torch::jit::script::Module类型的变量用于加载模型

  torch::jit::script::Module m_lumaModuleISlice[3];
  torch::jit::script::Module m_chromaModuleISlice[3];
  torch::jit::script::Module m_lumaModuleBSlice[5];
  torch::jit::script::Module m_chromaModuleBSlice[5];

 通过torch::jit::load函数加载模型

void CNNFilter::initISlice(int qp)
{
  if (m_modelInitFlagISlice)
    return;
  
  at::set_num_threads(1);
  at::set_num_interop_threads(1);

  for (int i = 0; i < 3; i++)
  {
    std::string sLumaModelName = "E:/VTM/models/JVET_W_EE_LumaCNNFilter_IntraSlice_qp" + sQp[i] + ".pt";
    std::string sChromaModelName = "E:/VTM/models/JVET_W_EE_ChromaCNNFilter_IntraSlice_qp" + sQp[i] + ".pt";
    m_lumaModuleISlice[i] = torch::jit::load(sLumaModelName); // 加载模型
    m_chromaModuleISlice[i] = torch::jit::load(sChromaModelName);
  }

}

2. 创建模型的输入

 torch::Tensor张量类型，libtorch和pytorch中许多api一致，以产生全1张量为例，通过torch::ones({1, 1, blockSizeVer, blockSizeHor})产生尺寸为{1, 1, blockSizeVer, blockSizeHor}的四维张量，注意libtorch是使用{}表示尺寸。

定义好张量以后，可以通过调用data_ptr()函数获得指针，通过指针直接修改数据。

  torch::NoGradGuard no_grad_guard用于不计算梯度，和pytorch中的with torch.no_grad()等效

  torch::Tensor imageBatch = torch::ones({1, 1, blockSizeVer, blockSizeHor});
  float *pImageBatch = imageBatch.data_ptr<float>(); // 指向imageBatch的指针

  torch::Tensor predBatch = torch::ones({1, 1, blockSizeVer, blockSizeHor});
  float *pPredBatch = predBatch.data_ptr<float>(); // 指向predBatch 的指针

  // Create a vector of inputs. 创建vector用于模型输入
  std::vector<torch::jit::IValue> input; 
  input.push_back(imageBatch);
  input.push_back(predBatch);

  int idx = 0;
  int blockSize = blockSizeVer * blockSizeHor;
  torch::NoGradGuard no_grad_guard;
  torch::globalContext().setFlushDenormal(true);

  //at::init_num_threads();
  for (int yy = 0; yy < blockSizeVer; yy++)
  {
    for (int xx = 0; xx < blockSizeHor; xx++)
    {
      // 通过指针给tensor赋值
      idx = yy * strideRec + xx;
      pImageBatch[yy*blockSizeHor+xx] = pRec[idx] / maxValue;  

      idx = yy * stridePred + xx;
      pPredBatch[yy*blockSizeHor+xx] = pPred[idx] / maxValue;
    }
  }

3. 模型推理

module->forward(): 模型前向传播的函数，输入值建议使用vector类型

  // Execute the model and turn its output into a tensor. 推理模型并将输出转换为tensor
  at::Tensor output = m_lumaModuleISlice[modelIdx].forward(input).toTensor();
  float *pOutput = output.data_ptr<float>(); // 输出tensor的指针

模型返回值

如果模型只有一个返回值，那么常用如下语句toTensor转换成张量

at::Tensor result = module.forward({tensor_image}).toTensor();

如果模型有多个返回值，那么则需要转换成Tuple.

auto result = module.forward({tensor_image}).toTuple();
at::Tensor loc= result->elements()[0].toTensor();
at::Tensor conf = result->elements()[1].toTensor();