LibTorch Practice 2: MNIST’s libtorch code

Table of Contents

I. Introduction

2. Another way to download the data set

3. MNIST’s Pytorch source code

4. MNIST’s Libtorch source code

1. Foreword

We have introduced MNIST’s python training code and torchscript-based model serialization (export model) before. Today let’s see how to use libtorch C++ to implement handwritten digit training.

2. Another way to download the data set

At the same time, I have already said how to download your MNIST data set. Regarding the download of the data set, this unimportant issue has been stuck for a long time. It is a waste of time. Bad review. Here is another download method. In the official warehouse, there is a script that can directly download https://github.com/pytorch/examples/blob/main/cpp/tools/download_mnist.py and execute it directly in the command line window. It can be downloaded as follows. The network may be stuck, but it is downloaded.

Just post the download_mnist.py source code here:

from __future__ import division
from __future__ import print_function

import argparse
import gzip
import os
importsys
import urllib

try:
    from urllib.error import URLError
    from urllib.request import urlretrieve
except ImportError:
    from urllib2 import URLError
    from urllib import urlretrieve

RESOURCES = [
    'train-images-idx3-ubyte.gz',
    'train-labels-idx1-ubyte.gz',
    't10k-images-idx3-ubyte.gz',
    't10k-labels-idx1-ubyte.gz',
]


def report_download_progress(chunk_number, chunk_size, file_size):
    if file_size != -1:
        percent = min(1, (chunk_number * chunk_size) / file_size)
        bar = '#' * int(64 * percent)
        sys.stdout.write('\r0% |{:<64}| {}%'.format(bar, int(percent * 100)))


def download(destination_path, url, quiet):
    if os.path.exists(destination_path):
        if not quiet:
            print('{} already exists, skipping ...'.format(destination_path))
    else:
        print('Downloading {} ...'.format(url))
        try:
            hook = None if quiet else report_download_progress
            urlretrieve(url, destination_path, reporthook=hook)
        except URLError:
            raise RuntimeError('Error downloading resource!')
        finally:
            if not quiet:
                #Just a newline.
                print()


def unzip(zipped_path, quiet):
    unzipped_path = os.path.splitext(zipped_path)[0]
    if os.path.exists(unzipped_path):
        if not quiet:
            print('{} already exists, skipping ... '.format(unzipped_path))
        return
    with gzip.open(zipped_path, 'rb') as zipped_file:
        with open(unzipped_path, 'wb') as unzipped_file:
            unzipped_file.write(zipped_file.read())
            if not quiet:
                print('Unzipped {} ...'.format(zipped_path))


def main():
    parser = argparse.ArgumentParser(
        description='Download the MNIST dataset from the internet')
    parser.add_argument(
        '-d', '--destination', default='.', help='Destination directory')
    parser.add_argument(
        '-q',
        '--quiet',
        action='store_true',
        help="Don't report about progress")
    options = parser.parse_args()

    if not os.path.exists(options.destination):
        os.makedirs(options.destination)

    try:
        for resources in RESOURCES:
            path = os.path.join(options.destination, resource)
            url = 'http://yann.lecun.com/exdb/mnist/{}'.format(resource)
            download(path, url, options.quiet)
            unzip(path, options.quiet)
    except KeyboardInterrupt:
        print('Interrupted')


if __name__ == '__main__':
    main()

During the download process, it may be stuck. The download information is as follows:

(base) C:\Users\Administrator\Desktop\examples-master_2\examples-master\cpp\tools>python download_mnist.py
.\train-images-idx3-ubyte.gz already exists, skipping ...
.\train-images-idx3-ubyte already exists, skipping ...
.\train-labels-idx1-ubyte.gz already exists, skipping ...
.\train-labels-idx1-ubyte already exists, skipping ...
.\t10k-images-idx3-ubyte.gz already exists, skipping ...
.\t10k-images-idx3-ubyte already exists, skipping ...
.\t10k-labels-idx1-ubyte.gz already exists, skipping ...
.\t10k-labels-idx1-ubyte already exists, skipping ... 

Python code trains 5 epoch results.

Test set: Average loss: 0.0287, Accuracy: 9907/10000 (99%)

3. MNIST’s Pytorch source code

Python source code of MNIST:

from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR


class Net(nn.Module):
    def __init__(self): # self refers to the class instance object itself (note: not the class itself).
    # self is not a keyword
        # super is used for inheritance, https://www.runoob.com/python/python-func-super.html
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # input:28*28
        x = self.conv1(x) # -> (28 - 3 + 1 = 26),26*26*32
        x = F.relu(x)
        # input:26*26*32
        x = self.conv2(x) # -> (26 - 3 + 1 = 24),24*24*64
        # input:24*24*64
        x = F.relu(x)
        x = F.max_pool2d(x, 2)# -> 12*12*64 = 9216
        x = self.dropout1(x) #Do not change dimensions
        x = torch.flatten(x, 1) # 9216*1
        # w = 128*9216
        x = self.fc1(x) # -> 128*1
        x = F.relu(x)
        x = self.dropout2(x)
        # w = 10*128
        x = self.fc2(x) # -> 10*1
        output = F.log_softmax(x, dim=1) # softmax normalization
        return output


def train(args, model, device, train_loader, optimizer, epoch):
    # When using pytorch to build a neural network, a model.train() sentence will be added above the program during the training process.
    #The function is to enable batch normalization and drop out.
    # Model.eval() will be used during the test. At this time, the neural network will use the batch normalization value and will not use drop out.
    model.train()
    # You can view the parameter size of the convolution kernel
    #model.conv1.weight.shape torch.Size([32, 1, 3, 3]
    #model.conv2.weight.shape torch.Size([64, 32, 3, 3])

    for batch_idx, (data, target) in enumerate(train_loader):
        # train_loader.dataset.data.shape
        # Out[9]: torch.Size([60000, 28, 28])

        # batch_size:64
        # data: 64 sample input, torch.Size([64, 1, 28, 28])
        # target: 64 labels,torch.Size([64])
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        # output:torch.Size([64, 10])
        output = model(data)
        # Similar to cross entropy
        # reference: https://blog.csdn.net/qq_22210253/article/details/85229988
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        # Let’s print a convolution kernel parameter to see
        # print(model.conv2._parameters)

        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
            if args.dry_run:
                break


def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss + = F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
            correct + = pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\
Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\
'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))


def main():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=5, metavar='N',
                        help='number of epochs to train (default: 14)')
    parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
                        help='learning rate (default: 1.0)')
    parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
                        help='Learning rate step gamma (default: 0.7)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--dry-run', action='store_true', default=False,
                        help='quickly check a single pass')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')
    parser.add_argument('--save-model', action='store_true', default=True,
                        help='For Saving the current Model')
    args = parser.parse_args()
    use_cuda = not args.no_cuda and torch.cuda.is_available()

    torch.manual_seed(args.seed)

    device = torch.device("cuda" if use_cuda else "cpu")

    train_kwargs = {'batch_size': args.batch_size}
    test_kwargs = {'batch_size': args.test_batch_size}
    if use_cuda:
        cuda_kwargs = {'num_workers': 1,
                       'pin_memory': True, # Lock page memory, which can speed up memory to video memory
                       'shuffle': True}
        train_kwargs.update(cuda_kwargs)
        test_kwargs.update(cuda_kwargs)
    # torchvision.transforms is an image preprocessing package in pytorch. Compose is generally used to integrate multiple steps together.
    #
    transform = transforms.Compose([
        transforms.ToTensor(), # (H x W x C), [0, 255] -> (C x H x W), [0.0, 1.0]
        transforms.Normalize((0.1307,), (0.3081,)) # Normalization of data
        ])

    dataset1 = datasets.MNIST('../data', train=True, download=True,
                       transform=transform)
    dataset2 = datasets.MNIST('../data', train=False,
                       transform=transform)
    train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)

    model = Net().to(device)
    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
    # Fixed step attenuation
    # reference: https://zhuanlan.zhihu.com/p/93624972
    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)
        scheduler.step()

    if args.save_model:
        #torch.save(model.state_dict(), "pytorch_mnist.pt")
        torch.save(model, "pytorch_mnist.pth")


if __name__ == '__main__':
    main()

4. MNIST’s Libtorch source code

The following is the C++ code (the network results of the official C++ code do not seem to completely correspond to the python code, so I made modifications, which actually changed the network model, please see struct Net: torch::nn::Module ): You can compare struct Net: torch::nn::Module and class Net(nn.Module) in the above python code:

#include<torch/torch.h>
#include<cstddef>
#include<iostream>
#include<vector>
#include<string>
//Inherited from Module module
struct Net : torch::nn::Module
{
    // Constructor
    Net():
        conv1(torch::nn::Conv2dOptions(1, 32, 3)), // kernel_size = 5
        conv2(torch::nn::Conv2dOptions(32, 64, 3)),
        fc1(9216, 128),
        fc2(128, 10)
    {
        register_module("conv1", conv1);
        register_module("conv2", conv2);
        register_module("conv2_drop", conv2_drop);
        register_module("fc1", fc1);
        register_module("fc2", fc2);
    }
    //Member function: forward propagation
    torch::Tensor forward(torch::Tensor x)
    {
        // input:1*28*28
        x = torch::relu(conv1->forward(x)); //conv1:(28 - 3 + 1 = 26), 26*26*32
        // input:26*26*32
        x = torch::max_pool2d(torch::relu(conv2->forward(x)), 2);//conv2:(26 - 3 + 1 = 24),24*24*64; max_poolded:12*12* 64 = 9216
        x = torch::dropout(x, 0.25, is_training());
        x = x.view({ -1, 9216 });// 9216*1
        // w:128*9216
        x = torch::relu(fc1->forward(x)); //fc1:w = 128*9216,w * x ->128*1
        x = torch::dropout(x, 0.5, is_training());
        // w:10*128
        x = fc2->forward(x);//fc2:w = 10*128,w * x -> 10*1
        x = torch::log_softmax(x, 1);
        return x;

    }


    //Module members
    torch::nn::Conv2d conv1;
    torch::nn::Conv2d conv2;
    torch::nn::Dropout2d conv2_drop;
    torch::nn::Linear fc1;
    torch::nn::Linear fc2;
};

//train
template<typename DataLoader>
void train(size_t epoch, Net & amp; model, torch::Device device, DataLoader & amp; data_loader, torch::optim::Optimizer & amp; optimizer, size_t dataset_size)
{
    //set "train" mode
    model.train();
    size_t batch_idx = 0;
    for (auto & amp; batch: data_loader)
    {
        auto data = batch.data.to(device);
        auto 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();

        //Print loss every 10 batch_size
        if (batch_idx + + % 10 == 0)
        {
            std::printf("\rTrain Epoch: %ld [%5ld/%5ld] Loss: %.4f",
                epoch,
                batch_idx * batch.data.size(0),
                dataset_size,
                loss.template item<float>());
        }
    }
}

template<typename DataLoader>
void test(Net & amp; model, torch::Device device, DataLoader & amp; data_loader, size_t dataset_size)
{
    torch::NoGradGuard no_grad;
    // set "test" mode
    model.eval();
    double test_loss = 0;
    int32_t correct = 0;
    for (const auto & amp; batch: data_loader)
    {
        auto data = batch.data.to(device);
        auto targets = batch.target.to(device);
        auto output = model.forward(data);
        test_loss + = torch::nll_loss(output, targets, /*weight=*/{}, torch::Reduction::Sum).template item<float>();
        auto pred = output.argmax(1);
        // eq = equal determines whether prediction is equal to label
        correct + = pred.eq(targets).sum().template item<int64_t>();
    }
    test_loss /= dataset_size;
    std::printf(
        "\
Test set: Average loss: %.4f | Accuracy: %.3f\
",
        test_loss,
        static_cast<double>(correct) / dataset_size);
}

int main()
{
    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;
    model.to(device);
    // load train data
    auto train_dataset = torch::data::datasets::MNIST("D://MNIST//")
        .map(torch::data::transforms::Normalize<>(0.1307, 0.3081))
        .map(torch::data::transforms::Stack<>());

    const size_t train_dataset_size = train_dataset.size().value();
    std::cout << train_dataset_size << std::endl;
    auto train_loader = torch::data::make_data_loader<torch::data::samplers::SequentialSampler>(
        std::move(train_dataset), 64);
    // load test data
    auto test_dataset = torch::data::datasets::MNIST(
        "D://MNIST//", 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), 1000);

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

    //train
    for (size_t epoch = 0; epoch < 5; epoch + + )
    {
        train(epoch, model, device, *train_loader, optimizer, train_dataset_size);
        test(model, device, *test_loader, test_dataset_size);
    }
    //save
    return 1;
}

The C++ code training results are as follows:

You can see that the C++ version of the MNIST code can train the model normally