libtorch c++ 线性卷积联合网络的训练及测试用于识别MNIST手写数据集

本实例同时采用卷积、池化、丢弃、非线性化、和线性网络层等多种网格联合识别手写数字。

（1）网络定义模块

网络的定义部分，定义结构体Net，内部成员有二维卷积层conv1，conv2，丢弃层conv2_dropout，线性层fc1，fc2，采用的其它网络层有最大池化层，max_pool2d, 非线性化层relu，压缩值域层log_softmax。

其中，卷积层和线性层的参数如下：

conv1(torch::nn::Conv2dOptions(1, 10, /kernel_size=/5)), conv2(torch::nn::Conv2dOptions(10, 20, /kernel_size=/5)), fc1(320, 50), fc2(50, 10)

丢弃层的参数是torch::dropout(x, 0.5)，意味着随机的选择一半的数据被设置为0

池化层的参数是：

torch::max_pool2d(x, 2)

非线性化层采用的是默认函数torch::relu(x)

压缩值域层的参数是torch::log_softmax(x, 1)，它的输出是对数概率

整个的数据处理流程是：

data-->conv1-->max_pool2d-->relu-->conv2-->conv2_drop-->max_pool2d-->view-->fc1-->relu-->dropout-->fc2

其中可以引起数据尺寸变化的层是conv1, conv2, max_pool2d, view, fc1, fc2，数据在各个网络层传递过程中，尺寸变化是：

struct Net : torch::nn::Module {
    Net()
        : conv1(torch::nn::Conv2dOptions(1, 10, /*kernel_size=*/5)),
        conv2(torch::nn::Conv2dOptions(10, 20, /*kernel_size=*/5)),
        fc1(320, 50),
        fc2(50, 10) {
        register_module("conv1", conv1);
        register_module("conv2", conv2);
        register_module("conv2_drop", conv2_drop);
        register_module("fc1", fc1);
        register_module("fc2", fc2);
    }

    torch::Tensor forward(torch::Tensor x) {
        x = torch::relu(torch::max_pool2d(conv1->forward(x), 2));
        x = torch::relu(
            torch::max_pool2d(conv2_drop->forward(conv2->forward(x)), 2));
        x = x.view({ -1, 320 });
        x = torch::relu(fc1->forward(x));
        x = torch::dropout(x, /*p=*/0.5, /*training=*/is_training());
        x = fc2->forward(x);
        return torch::log_softmax(x, /*dim=*/1);
    }

    torch::nn::Conv2d conv1;
    torch::nn::Conv2d conv2;
    torch::nn::Dropout conv2_drop;
    torch::nn::Linear fc1;
    torch::nn::Linear fc2;
};
//TORCH_MODULE(Net); //to save and load  model 

（2）训练函数模块

训练函数train（）采用了模板的方式，方便函数的通用性，对于其他网络模型和数据加载器，都可以采用这个训练函数进行训练。

输入的数据：样本的世代epoch，神经网路模型model，模型所在的设备device，样本加载器data_loader，优化器optimizer，该世代样本集大小dataset_size，

训练模块首先通过model.train()函数，启动训练。

对每一代数据，从数据加载器中一批一批地加载样本，对每批样本batch，通过batch.data，获取每批样本的数据项data，通过batch.target获取样本的标签项target，然后就可以开始真正的模型训练，在训练之前需要先把优化器的梯度设置为0，optimizer.zero_grad()。下面是训练过程：首先，把数据项传递给网络执行网络模型的正向传播计算model->forward(data), 正向传播计算结果是output，然后，计算损失函数loss=torch::nill_loss(output, targets)，并开始反向传播计算梯度loss.backward()，最后，通过优化器对梯度进行优化optimizer.step()。除了训练过程本身，还需要每个一段输出训练的状态，主要是目标函数的值。训练结束后是对模型的保存，torch::save(model, "model.pt")

template <class T, typename DataLoader>
void train(
    size_t epoch,
    T& model,
    torch::Device device,
    DataLoader& data_loader,
    torch::optim::Optimizer& optimizer,
    size_t dataset_size) {
    model->train();
    size_t batch_idx = 0;
    for (auto& batch : data_loader) {
        auto data = batch.data.to(device), targets = batch.target.to(device);
        optimizer.zero_grad();
        auto output = model->forward(data);
        auto loss = torch::nll_loss(output, targets);
        AT_ASSERT(!std::isnan(loss.template item<float>()));
        loss.backward();       
        optimizer.step();

        if (batch_idx++ % kLogInterval == 0) {
            std::printf(
                "\rTrain Epoch: %ld [%5ld/%5ld] Loss: %.4f",
                epoch,
                batch_idx * batch.data.size(0),
                dataset_size,
                loss.template item<float>());
            std::ofstream sw("loss.txt", std::ios::app);
            int seq = batch_idx * batch.data.size(0) + (epoch -1) * dataset_size;
            sw << epoch <<"\t"<< batch_idx <<"\t"<< seq << "\t" << loss.template item<float>() << std::endl;
            sw.close();

        }
    }
}

（3）训练后模型的测试模块

首先加载已训练的模型，torch::load(model, "model.pt")

测试函数test（）同样采用了模板的方式，可以接收不同类型的网络和数据加载器

测试函数的输入是：神经网络模型model，设备device，样本加载器data_loader, 该世代样本集的大小dataset_size。

然后启动神经网络模型的评价函数，model.eval(),

对该世代中的每批数据batch，通过batch.data和batch.target获取数据项data和样本标签项target，通过模型的正向传播model->forward(data)得到计算结果output，通过torch::nll_loss()计算得到误差，nll_loss是负对数似然损失函数，它的输入是两个张量，同时还有一些选项参数，比如权重weight，降维Reduction方式：比如不降维None，均值Mean，求和Sum等。

除了计算误差值，还需要计算吻合率。正演传播的输出output是长度为10的向量，该向量包含0~9取值的概率，argmax（output）的作用是获取概率最大的那个预测值pred，在判断预测的pred与targets是否相等，对相等那部分求和，占总样本集的比例，即为吻合率。

template <class T, typename DataLoader>
void test(
    T& model,
    torch::Device device,
    DataLoader& data_loader,
    size_t dataset_size) {
    torch::NoGradGuard no_grad;
    model->eval();
    double test_loss = 0;
    int32_t correct = 0;
    for (const auto& batch : data_loader) {
        auto data = batch.data.to(device), targets = batch.target.to(device);
        auto output = model->forward(data);
        test_loss += torch::nll_loss(
            output,
            targets,
            /*weight=*/{},
            at::Reduction::Sum)
            .template item<float>();
        auto pred = output.argmax(1);
        correct += pred.eq(targets).sum().template item<int64_t>();
    }

    test_loss /= dataset_size;
    std::printf(
        "\nTest set: Average loss: %.4f | Accuracy: %.3f\n",
        test_loss,
        static_cast<double>(correct) / dataset_size);
}

（4）数据集定义与数据加载器设置

通过torch::data::dataset名称空间下的MNIST类型加载MNIST数据集，MNIST的构造函数是MNIST(const std::string& root, Mode mode = Mode::kTrain);只需指定MNIST数据所在的路径即可，样本的默认用图是训练，Mode::KTrain，可以不设置；如果数据集的目的是测试，则需要设置Mode::KTest，

数据加载后一般需要通过data下的转换函数的正态化，torch::data::transform::Normalize对数据进行正态化（指定均质和方差），然后通过torch::data::transform::Stack把该批数据叠置/整合一个tensor.

数据集定义之后还要定义数据加载器data_loader,方法是torch::data::make_data_loader，它是个模板函数，可以建立不同类型的数据加载器，比如序列式torch::data::samplers::SequentialSampler的采样器，随机化的采样器torch::data::samplers::RandomSampler，以及两者分布式的版本，然后的参数是是否移动数据，以及每个批次样本的个数。

类似的方法可以创建测试数据集和测试数据的加载器，

auto train_dataset = torch::data::datasets::MNIST(kDataRoot)
        .map(torch::data::transforms::Normalize<>(0.1307, 0.3081))
        .map(torch::data::transforms::Stack<>());
    const size_t train_dataset_size = train_dataset.size().value();
    auto train_loader =
        torch::data::make_data_loader<torch::data::samplers::RandomSampler::SequentialSampler>(
            std::move(train_dataset), kTrainBatchSize);

    auto test_dataset = torch::data::datasets::MNIST(
        kDataRoot, torch::data::datasets::MNIST::Mode::kTest)
        .map(torch::data::transforms::Normalize<>(0.1307, 0.3081))
        .map(torch::data::transforms::Stack<>());
    const size_t test_dataset_size = test_dataset.size().value();
    auto test_loader =
        torch::data::make_data_loader(std::move(test_dataset), kTestBatchSize);

（5）主控函数

在主控函数中，（1）定义一些常用用于设置数据集、数据加载器等，（2）新建网络模型的实例，（3）创建训练数据集、数据加载器和测试数据集和加载器，（4）定义优化器，（5）通过训练训练和测试模型，（6）训练模型的保存和加载，（7）测试最后的模型效果。

auto main() -> int {
    torch::manual_seed(1);

    torch::DeviceType device_type;
    if (torch::cuda::is_available()) {
        std::cout << "CUDA available! Training on GPU." << std::endl;
        device_type = torch::kCUDA;
    }
    else {
        std::cout << "Training on CPU." << std::endl;
        device_type = torch::kCPU;
    }
    torch::Device device(device_type);

    //    Net model;

    auto model = std::make_shared<Net>();
    model->to(device);

    auto train_dataset = torch::data::datasets::MNIST(kDataRoot)
        .map(torch::data::transforms::Normalize<>(0.1307, 0.3081))
        .map(torch::data::transforms::Stack<>());
    const size_t train_dataset_size = train_dataset.size().value();
    auto train_loader =
        torch::data::make_data_loader<torch::data::samplers::RandomSampler::SequentialSampler>(
            std::move(train_dataset), kTrainBatchSize);

    auto test_dataset = torch::data::datasets::MNIST(
        kDataRoot, torch::data::datasets::MNIST::Mode::kTest)
        .map(torch::data::transforms::Normalize<>(0.1307, 0.3081))
        .map(torch::data::transforms::Stack<>());
    const size_t test_dataset_size = test_dataset.size().value();
    auto test_loader =
        torch::data::make_data_loader(std::move(test_dataset), kTestBatchSize);

    torch::optim::SGD optimizer(
        model->parameters(), torch::optim::SGDOptions(0.01).momentum(0.5));

    for (size_t epoch = 1; epoch <= kNumberOfEpochs; ++epoch) {
        train(epoch, model, device, *train_loader, optimizer, train_dataset_size);
        test(model, device, *test_loader, test_dataset_size);
    }

    torch::save(model, "model.pt");//save model
    torch::load(model, "model.pt");//load model

    //print model parameters
    for (const auto& pair : model->named_parameters()) {
        std::cout << pair.key() << ": " << pair.value().sizes() << std::endl;
    }
    test(model, device, *test_loader, test_dataset_size);

}