libtorch教程（四）

Posted on | | Visitors: | 本文总阅读量 次

本章将详细介绍如何使用libtorch自带的数据加载模块，使用该模块是实现模型训练的重要条件。除非这个数据加载模块功能不够，不然继承libtorch的数据加载类还是很有必要的，简单高效。

使用前置条件

libtorch提供了丰富的基类供用户自定义派生类，torch::data::Dataset就是其中一个常用基类。使用该类需要明白基类和派生类，以及所谓的继承和多态。有c++编程经验者应该都不会陌生，为方便不同阶段读者就简单解释一下吧。类就是父亲，可以生出不同的儿子，生儿子叫派生或者继承(看使用语境)，生不同的儿子就实现了多态。父亲就是基类，儿子就是派生类。现实中，父亲会把自身的一部分财产留下来养老，儿子们都不能碰，这就是private了，部分财产儿子能用，但是儿子的对象不能用，这叫protected，还有些财产谁都能用就是public。和现实中的父子类似，代码中，派生类可以使用父类的部分属性或者函数，全看父类怎样定义。

然后理解一下虚函数，就是父亲指定了部分财产是public的，但是是用来买房的，不同的儿子可以买不同的房子，可以全款可以贷款，这就是财产在父亲那就是virtual的。子类要继承这个virtual财产可以自己重新规划使用方式。

事实上，如果有过pytorch的编程经验者很快会发现，libtorch的Dataset类的使用和python下使用非常相像。pytorch自定义dataload，需要定义好Dataset的派生类，包括初始化函数__init__，获取函数__getitem__以及数据集大小函数__len__。类似的，libtorch中同样需要处理好初始化函数，get()函数和size()函数。

图片文件遍历

下面以分类任务为例，介绍libtorch的Dataset类的使用。使用pytorch官网提供的昆虫分类数据集，下载到本地解压。将该数据集根目录作为索引，实现Dataloader对图片的加载。

首先定义一个加载图片的函数，使用网上出现较多的c++遍历文件夹的代码，将代码稍作修改如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
//遍历该目录下的.jpg图片
void load_data_from_folder(std::string image_dir, std::string type, std::vector<std::string> &list_images, std::vector<int> &list_labels, int label);

void load_data_from_folder(std::string path, std::string type, std::vector<std::string> &list_images, std::vector<int> &list_labels, int label)
{
    long long hFile = 0; //句柄
    struct _finddata_t fileInfo;
    std::string pathName;
    if ((hFile = _findfirst(pathName.assign(path).append("\\*.*").c_str(), &fileInfo)) == -1)
    {
        return;
    }
    do
    {
        const char* s = fileInfo.name;
        const char* t = type.data();

        if (fileInfo.attrib&_A_SUBDIR) //是子文件夹
        {
            //遍历子文件夹中的文件(夹)
            if (strcmp(s, ".") == 0 || strcmp(s, "..") == 0) //子文件夹目录是.或者..
                continue;
            std::string sub_path = path + "\\" + fileInfo.name;
            label++;
            load_data_from_folder(sub_path, type, list_images, list_labels, label);

        }
        else //判断是不是后缀为type文件
        {
            if (strstr(s, t))
            {
                std::string image_path = path + "\\" + fileInfo.name;
                list_images.push_back(image_path);
                list_labels.push_back(label);
            }
        }
    } while (_findnext(hFile, &fileInfo) == 0);
    return;
}

修改后的函数接受数据集文件夹路径image_dir和图片类型image_type，将遍历到的图片路径和其类别分别存储到list_images和list_labels，最后lable变量用于表示类别计数。传入lable=-1，返回的lable值加一后等于图片类别。

自定义Dataset

定义dataSetClc，该类继承自torch::data::Dataset。定义私有变量image_paths和labels分别存储图片路径和类别，是两个vector变量。dataSetClc的初始化函数就是加载图片和类别。通过get()函数返回由图像和类别构成的张量列表。可以在get()函数中做任意针对图像的操作，如数据增强等。效果等价于pytorch中的__getitem__中的数据增强。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class dataSetClc:public torch::data::Dataset<dataSetClc>{
public:
    int class_index = 0;
    dataSetClc(std::string image_dir, std::string type){
        load_data_from_folder(image_dir, std::string(type), image_paths, labels, class_index-1);
    }
    // Override get() function to return tensor at location index
    torch::data::Example<> get(size_t index) override{
        std::string image_path = image_paths.at(index);
        cv::Mat image = cv::imread(image_path);
        cv::resize(image, image, cv::Size(224, 224)); //尺寸统一，用于张量stack，否则不能使用stack
        int label = labels.at(index);
        torch::Tensor img_tensor = torch::from_blob(image.data, { image.rows, image.cols, 3 }, torch::kByte).permute({ 2, 0, 1 }); // Channels x Height x Width
        torch::Tensor label_tensor = torch::full({ 1 }, label);
        return {img_tensor.clone(), label_tensor.clone()};
    }
    // Override size() function, return the length of data
    torch::optional<size_t> size() const override {
        return image_paths.size();
    };
private:
    std::vector<std::string> image_paths;
    std::vector<int> labels;
};

使用自定义的Dataset

下面使用定义好的数据加载类，以昆虫分类中的训练集作为测试，代码如下。可以打印加载的图片张量和类别。

1
2
3
4
5
6
7
8
9
int batch_size = 2;
std::string image_dir = "your path to\\hymenoptera_data\\train";
auto mdataset = myDataset(image_dir,".jpg").map(torch::data::transforms::Stack<>());
auto mdataloader = torch::data::make_data_loader<torch::data::samplers::RandomSampler>(std::move(mdataset), batch_size);
for(auto &batch: *mdataloader){
    auto data = batch.data;
    auto target = batch.target;
    std::cout<<data.sizes()<<target;
}

代码见github