How to build a dataset using PyTorch C++ frontend

smth · March 13, 2019, 8:50am

Creating a C++ dataset is similar to creating a PyTorch python dataset.

One has to subclass the Dataset class, and then implement some base methods.

See the MNIST class for example:

pytorch/pytorch/blob/8311bbee7f1dd33346f18c769cfeb8c5b5941874/torch/csrc/api/include/torch/data/datasets/mnist.h

#pragma once

#include <torch/data/datasets/base.h>
#include <torch/data/example.h>
#include <torch/types.h>

#include <torch/csrc/WindowsTorchApiMacro.h>

#include <cstddef>
#include <string>

namespace torch {
namespace data {
namespace datasets {
/// The MNIST dataset.
class TORCH_API MNIST : public Dataset<MNIST> {
 public:
  /// The mode in which the dataset is loaded.
  enum class Mode { kTrain, kTest };

This file has been truncated. show original

github.com

pytorch/pytorch/blob/8311bbee7f1dd33346f18c769cfeb8c5b5941874/torch/csrc/api/src/data/datasets/mnist.cpp

#include <torch/data/datasets/mnist.h>

#include <torch/data/example.h>
#include <torch/types.h>

#include <c10/util/Exception.h>

#include <cstddef>
#include <fstream>
#include <string>
#include <vector>

namespace torch {
namespace data {
namespace datasets {
namespace {
constexpr uint32_t kTrainSize = 60000;
constexpr uint32_t kTestSize = 10000;
constexpr uint32_t kImageMagicNumber = 2051;
constexpr uint32_t kTargetMagicNumber = 2049;

This file has been truncated. show original

Alternatively, you can subclass from BatchDataset and implement a get_batch method that returns a batch of samples directly: https://github.com/pytorch/pytorch/blob/master/torch/csrc/api/include/torch/data/datasets/base.h#L40-L94