libtorch c++ 搭建分类网络进行训练和预测
目录
1. vgg.h
2. vgg.cpp
3. 预训练权重
3.1 保存pytorch版的预训练权重
3.2 训练
3.2.1 打印权重参数信息
3.2.2 权重初始化
3.2.3 测试预训练权重是否可用
3.2.4 训练函数
3.2.5 主函数
4. 测试
4.1 主函数
4.2 predic函数
1. vgg.h
这里以vgg分类网络为例，vgg网络简单：多层卷积提取特征，最大池化下采样，全连接层进行分类，最后接一个sotmax归一化输出。
#ifndef VGG_H
#define VGG_H
#include
#include
#include
// 二维卷积参数配置
// 注意：关键字inline必须与函数定义放在一起才能使函数成为内联函数，仅仅将inline放在函数声明前面不起任何作用。
inline torch::nn::Conv2dOptions conv_options(int64_t in_planes, int64_t out_planes, int64_t kerner_size,
int64_t stride = 1, int64_t padding = 0, bool with_bias = false) {
torch::nn::Conv2dOptions conv_options = torch::nn::Conv2dOptions(in_planes, out_planes, kerner_size);
conv_options.stride(stride);
conv_options.padding(padding);
conv_options.bias(with_bias);
return conv_options;
}
// 最大池化参数配置
inline torch::nn::MaxPool2dOptions maxpool_options(int kernel_size, int stride){
torch::nn::MaxPool2dOptions maxpool_options(kernel_size);
maxpool_options.stride(stride);
return maxpool_options;
}
// 工具函数：生成多个卷积层，作为特征提取层
// 参数cfg: 中间特征层的通道数，batch_norm：是否使用bn
torch::nn::Sequential make_features(std::vector &cfg, bool batch_norm);
// 创建VGG类，必须继承Module类
// vgg网络：卷积层+pooling层+全连接层
class VGGImpl: public torch::nn::Module
{
public:
VGGImpl(std::vector& cfg, int num_classes = 1000, bool batch_norm = false); // 1个构造函数，初始化各个层操作的参数
torch::Tensor forward(torch::Tensor x); // 1个forward函数
private:
torch::nn::Sequential features_{nullptr}; // 卷积层
torch::nn::AdaptiveAvgPool2d avgpool{nullptr}; // pooling
torch::nn::Sequential classifier; // 全连接层
};
TORCH_MODULE(VGG);
#endif // VGG_H
2. vgg.cpp
#include "vgg.h"
// 工具函数：生成多个卷积层，作为特征提取层
// 参数cfg: 中间特征层的通道数，batch_norm：是否使用bn
torch::nn::Sequential make_features(std::vector &cfg, bool batch_norm){
torch::nn::Sequential features;
int in_channels = 3;
for(auto v : cfg){ // v是通道数
if(v==-1){ // 遇到-1，则接一个池化层
features->push_back(torch::nn::MaxPool2d(maxpool_options(2,2)));
}
else{
// conv2 + bn + relu
auto conv2d = torch::nn::Conv2d(conv_options(in_channels,v,3,1,1)); // k=3,s=1,p=1
features->push_back(conv2d);
if(batch_norm){
features->push_back(torch::nn::BatchNorm2d(torch::nn::BatchNorm2dOptions(v))); // or torch::nn::BatchNorm2d(v)
}
features->push_back(torch::nn::ReLU(torch::nn::ReLUOptions(true))); // relu层
in_channels = v;
}
}
return features;
}
// 初始化私有成员: features_，avgpool，classifier
VGGImpl::VGGImpl(std::vector &cfg, int num_classes, bool batch_norm){
features_ = make_features(cfg,batch_norm); // 1.初始化卷积层（包含了池化层）
avgpool = torch::nn::AdaptiveAvgPool2d(torch::nn::AdaptiveAvgPool2dOptions(7)); // 2.初始化池化层
classifier->push_back(torch::nn::Linear(torch::nn::LinearOptions(512 * 7 * 7, 4096))); // 3.初始化全连接层
classifier->push_back(torch::nn::ReLU(torch::nn::ReLUOptions(true)));
classifier->push_back(torch::nn::Dropout());
classifier->push_back(torch::nn::Linear(torch::nn::LinearOptions(4096, 4096)));
classifier->push_back(torch::nn::ReLU(torch::nn::ReLUOptions(true)));
classifier->push_back(torch::nn::Dropout());
classifier->push_back(torch::nn::Linear(torch::nn::LinearOptions(4096, num_classes))); // 两个linear+relu+dropout, 再加一个linear.
features_ = register_module("features",features_); // module必须注册
classifier = register_module("classifier",classifier);
}
torch::Tensor VGGImpl::forward(torch::Tensor x){
x = features_->forward(x); // 1，先卷积（包含了池化层），注册的module需要通过forward使用
x = avgpool(x); // 2，再池化，直接使用
x = torch::flatten(x,1); // 3，二维变一维
x = classifier->forward(x); // 4，再接全连接层
return torch::log_softmax(x, 1); // 5，最后softmax
}
3. 预训练权重
使用预训练权重训练模型，可以更快的收敛。为了使c++训练模型时能够使用pytorch版的预训练权重，c++代码搭建的vgg必须和pytorch一致。感兴趣的可以看下官方的pytorch代码torchvision.models.vgg16_bn
3.1 保存pytorch版的预训练权重
注意要使用jit.trace函数。
import torch
from torchvision.models import vgg16, vgg16_bn
# 在c++中搭建一个和pytorch下完全一致的vgg16bn。如果不一致的话其实不影响正常的模型训练和预测，
# 但是影响初始化状态，模型加载从ImageNet数据集训练好的权重以后，训练收敛的速度和收敛后的精度都会好很多。
model = vgg16_bn(pretrained=True)
model = model.to(torch.device("cpu"))
model.eval()
var = torch.ones((1, 3, 224, 224))
# 保存pytorch模型的权重不能直接用torch.save保存模型，这样存下来的模型不能被c++加载。我们利用部署时常用的torch.jit.script模型来保存。
traced_script_module = torch.jit.trace(model, var)
traced_script_module.save("vgg16bn.pt")
# 这样，模型的卷积层，归一化层，线性层的权重就保存到.pt文件中了。
3.2 训练
基于VGG类，再打包一次得到Classifier类。
classification.h
#ifndef CLASSIFICATION_H
#define CLASSIFICATION_H
#include
#include
#include
class Classifier
{
private:
torch::Device device = torch::Device(torch::kCPU); // 默认使用cpu
VGG vgg = VGG{nullptr}; // 自定义网络对象
public:
Classifier(int gpu_id = 0); // 构造函数，初始化device私用成员
void Initialize(int num_classes, std::string pretrained_path); // 加载预训练权重
void Train(int epochs, int batch_size, float learning_rate, std::string train_val_dir, std::string image_type, std::string save_path);
int Predict(cv::Mat &image); // preprocess + infer + postprocess
void LoadWeight(std::string weight); // 加载权重
};
#endif // CLASSIFICATION_H
3.2.1 打印权重参数信息
先看看权重参数的特点。
# pytorch
import torch
from torchvision.models import vgg16, vgg16_bn
model = vgg16_bn(pretrained=True)
for k, v in model.named_parameters():
print(k)
// c++
std::vector cfg_dd = {64, 64, -1, 128, 128, -1, 256, 256, 256, -1, 512, 512, 512, -1, 512, 512, 512, -1};
auto vgg_dd = VGG(cfg_dd,1000,true); // 直接实例化VGG即可
auto dictdd = vgg_dd->named_parameters();
for (auto n = dictdd.begin(); n != dictdd.end(); n++) // 打印出模型每一层（有权重的层，不包括类似激活函数层）的名称
{
std::cout<<(*n).key()< 左边是pytorch打印的信息，右边是c++打印的信息，可以发现c++的少了部分bias 权重，这是因为有些卷积没有设置bias（可以设置）。初始化的时候，根据权重参数名称初始化右边对应的值。
3.2.2 权重初始化
利用pytorch版的预训练权重（前面有保存下来），初始化c++网络权重参数。
classification.cpp
// 使用预训练权重，初始化自定义的模型权重
void Classifier::Initialize(int _num_classes, std::string _pretrained_path){
std::vector cfg_d = {64, 64, -1, 128, 128, -1, 256, 256, 256, -1, 512, 512, 512, -1, 512, 512, 512, -1};
auto net_pretrained = VGG(cfg_d,1000,true); // 注意这里是1000个类别，实例化预训练网络，为了取出权重
vgg = VGG(cfg_d,_num_classes,true); // 注意这里的类别个数
torch::load(net_pretrained, _pretrained_path); // 预训练的网络载入预训练的权重。
torch::OrderedDict pretrained_dict = net_pretrained->named_parameters();
torch::OrderedDict model_dict = vgg->named_parameters();
// 将训练的权重值提取出来，其中后面的分类层参数值丢弃不用。
for (auto n = pretrained_dict.begin(); n != pretrained_dict.end(); n++)
{
if (strstr((*n).key().data(), "classifier")) { // 不使用分类层的权重
continue;
}
model_dict[(*n).key()] = (*n).value();
}
torch::autograd::GradMode::set_enabled(false); // make parameters copying possible
auto new_params = model_dict; // implement this
auto params = vgg->named_parameters(true /*recurse*/);
auto buffers = vgg->named_buffers(true /*recurse*/);
for (auto& val : new_params) { // 功能：将屏蔽分类层后的新权重，复制到自定义的vgg网络模型权重中。
auto name = val.key();
auto* t = params.find(name); // 查看自定义网络中是否有该权重参数名
if (t != nullptr) {
t->copy_(val.value()); // 有，则将预训练得到的val复制到自定义网络参数中。
}
else {
t = buffers.find(name); // 没有，则看看buffer里面有没有，都没有则跳过。
if (t != nullptr) {
t->copy_(val.value());
}
}
}
torch::autograd::GradMode::set_enabled(true);
try
{
vgg->to(device);
}
catch (const std::exception&e)
{
std::cout << e.what() << std::endl;
}
return;
} 3.2.3 测试预训练权重是否可用
#include
int main(int argc, char *argv[])
{
//2,加载预训练权重，测试权重是否可用。
std::vector cfg_16bn = { 64, 64, -1, 128, 128, -1, 256, 256, 256, -1, 512, 512, 512, -1, 512, 512, 512, -1 };
auto vgg16bn = VGG(cfg_16bn, 1000, true);
torch::load(vgg16bn, "E:\\code\\c++\\LibtorchTutorials-main\\vgg16bn.pt");
vgg16bn->to(torch::Device(torch::kCUDA));
} 最后一句报错：PyTorch is not linked with support for cuda devices
解决办法：属性->链接器->命令行：
libtorch 1.5: -INCLUDE:THCudaCharTensor_zero
libtorch 1.6/1.7/1.9/1.9.1/1.10/1.1：-INCLUDE:?warp_size@cuda@at@@YAHXZ
libtorch 1.8.1: -INCLUDE:?wait@Future@ivalue@c10@@QEAAXXZ
或者：-INCLUDE:?mutate@OptOutMutator@cuda@fuser@jit@torch@@UEAAPEAVStatement@2345@PEAVForLoop@kir@2345@@Z3.2.4 训练函数
这里有个骚操作，前8个epoch保留前面的特征提取层参数不变（即使用预训练权重值），只更新分类层参数。
void Classifier::Train(int num_epochs, int batch_size, float learning_rate, std::string train_val_dir, std::string image_type, std::string save_path){
std::string path_train = train_val_dir+ "\\train"; // 数据路径
std::string path_val = train_val_dir + "\\val";
auto custom_dataset_train = dataSetClc(path_train, image_type).map(torch::data::transforms::Stack<>()); // dataset
auto custom_dataset_val = dataSetClc(path_val, image_type).map(torch::data::transforms::Stack<>());
auto data_loader_train = torch::data::make_data_loader(std::move(custom_dataset_train), batch_size); // dataloader
auto data_loader_val = torch::data::make_data_loader(std::move(custom_dataset_val), batch_size);
float loss_train = 0; float loss_val = 0; // 累加当前epoch内所有的loss，求平均loss
float acc_train = 0.0; float acc_val = 0.0; float best_acc = 0.0; // 累加准去率，和最佳准确率
for (size_t epoch = 1; epoch <= num_epochs; ++epoch) { // epoch
size_t batch_index_train = 0; // 记录当前是第几个batch
size_t batch_index_val = 0;
if (epoch == int(num_epochs / 2)) { learning_rate /= 10; } // 每两个epoch下降一次学习率
torch::optim::Adam optimizer(vgg->parameters(), learning_rate); // Learning Rate
if (epoch < int(num_epochs / 8)) // 前8个epoch只更新分类层参数
{
for (auto mm : vgg->named_parameters())
{
if (strstr(mm.key().data(), "classifier")) // 只更新分类层的参数
{
mm.value().set_requires_grad(true);
}
else
{
mm.value().set_requires_grad(false);
}
}
}
else { // 后面epoch次更新全部参数
for (auto mm : vgg->named_parameters())
{
mm.value().set_requires_grad(true);
}
}
// Iterate data loader to yield batches from the dataset
for (auto& batch : *data_loader_train) {
auto data = batch.data; // b,c,h,w
auto target = batch.target.squeeze(); // b,1
data = data.to(torch::kF32).to(device).div(255.0); // data: kf32,cuda,/255
target = target.to(torch::kInt64).to(device); // target: kint64, cuda
optimizer.zero_grad();
// Execute the model
torch::Tensor prediction = vgg->forward(data);
//cout << prediction << endl;
auto acc = prediction.argmax(1).eq(target).sum();
acc_train += acc.template item() / batch_size;
// Compute loss value
torch::Tensor loss = torch::nll_loss(prediction, target);
// Compute gradients
loss.backward();
// Update the parameters
optimizer.step();
loss_train += loss.item(); // 累加loss
batch_index_train++;
std::cout << "Epoch: " << epoch << " |Train Loss: " << loss.item() << " |Train Acc:" << acc_train / batch_index_train << "\r";
}
std::cout << std::endl;
//validation part
vgg->eval();
for (auto& batch : *data_loader_val) {
auto data = batch.data;
auto target = batch.target.squeeze();
data = data.to(torch::kF32).to(device).div(255.0);
target = target.to(torch::kInt64).to(device);
torch::Tensor prediction = vgg->forward(data);
// Compute loss value
torch::Tensor loss = torch::nll_loss(prediction, target);
auto acc = prediction.argmax(1).eq(target).sum(); // val准确率计算方法
acc_val += acc.template item() / batch_size;
loss_val += loss.item(); // 累加loss
batch_index_val++;
std::cout << "Epoch: " << epoch << " |Val Loss: " << loss_val / batch_index_val << " |Valid Acc:" << acc_val / batch_index_val << "\r";
}
std::cout << std::endl;
if (acc_val > best_acc) {
torch::save(vgg, save_path);
best_acc = acc_val;
}
loss_train = 0; loss_val = 0; acc_train = 0; acc_val = 0; batch_index_train = 0; batch_index_val = 0;
}
} 其中 dataset-dataSetClc的定义与上一篇博客一样都是图像分类数据集。
3.2.5 主函数
#include
int main(int argc, char *argv[])
{
//2，使用预训练权重，进行推理
std::string vgg_weight_path = "E:\\code\\c++\\LibtorchTutorials-main\\vgg16bn.pt";
std::string train_val_dir = "F:\\zxq\\data\\hymenoptera_data";
Classifier classifier(0); // gpu id: 0
classifier.Initialize(2, vgg_weight_path); // 使用预训练权重初始化权重，
classifier.Train(300, 4, 0.0003, train_val_dir, ".jpg", "classifer.pt");
} 验证集的准确率是0.91。
换成猫狗大战数据集（分出4000张训练），训练和测试集loss收敛的都很快。
4. 测试
测试的时候发现一个问题，在2080ti训练的模型，在3080显卡下预测会报错：
Microsoft C++ 异常: c10::Error，位于内存位置 0x000000301E6F3BF0 处。
可以在高端显卡训练，低端显卡部署。
4.1 主函数
// 3, infer
Classifier classifier(1); // gpu id: 0
classifier.Initialize(2, "E:\\code\\c++\\LibtorchTutorials-main\\vgg16bn.pt");
classifier.LoadWeight("E:\\code\\c++\\LibtorchLearning\\classifer.pt");
std::string train_val_dir = "F:\\zxq\\data\\custom\\dog";
cv::Mat image = cv::imread(train_val_dir+"\\dog.2344.jpg");
int cls_id = classifier.Predict(image); // preprocess + forward
std::cout << "cls_id: " << cls_id << std::endl; 
换成cat图片。
4.2 predic函数
int Classifier::Predict(cv::Mat& image){
// preprocess: resize, to_tensor, cuda, kf32, /255
cv::resize(image, image, cv::Size(448, 448));
torch::Tensor img_tensor = torch::from_blob(image.data, { image.rows, image.cols, 3 }, torch::kByte).permute({ 2, 0, 1 }); // c,h,w
img_tensor = img_tensor.to(device).unsqueeze(0).to(torch::kF32).div(255.0);
auto prediction = vgg->forward(img_tensor); // raw output
prediction = torch::softmax(prediction,1);
auto class_id = prediction.argmax(1);
std::cout<eval();
return;
} 参考：LibtorchTutorials/lesson5-TrainingVGG at main · AllentDan/LibtorchTutorials · GitHub