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);

}