跳转至

自定义 DataLoader

数据输入 一文所介绍,OneFlow 支持两种数据加载方式:直接使用 NumPy 数据或者使用 DataLoader 及其相关算子。

在大型工业场景下,数据加载容易成为训练的瓶颈。若使用 DataLoader 及相关预处理算子,因为有 OneFlow 内置的加速机制,可以更高效地加载和预处理数据,解决这个痛点。

在 OneFlow 中使用 DataLoader,一般通过调用 XXXReader 加载文件中的数据,调用 XXXDeocde 等对数据进行解码或其它预处理,他们一起协同完成 Dataloader 的功能。

OneFlow 目前内置了一些文件格式的 DataLoader。如果我们想使用 DataLoader 提高数据加载的效率,但是加载的数据格式暂时没有内置在 OneFlow 中,此时,我们可以自己实现 DataLoader,加载自定义的数据格式。

本文实现了一个 Mini Dataloader,在仓库中可查看完整代码

作为示例,Mini Dataloader 支持的文件格式为:以逗号分隔的每行两列数字的纯文本文件(见代码中的 part-000part-001 文件):

1.01,2.02
2.01,4.02
3.0,6.05
4.1,8.205
5,10
6.0,12.0
7.0,14.2
8.0,16.3
9.1,18.03

本文将以 Mini Dataloader 为例,对自定义格式的 DataLoader 的实现要点,进行讲解。

Dataloader 的组成

完整的 Dataloader 一般包括两类 Op:

  • Data Reader:负责将文件系统中的数据,加载到内存中的输入流,并最终将数据设置到 Op 的输出中
  • Data Decoder:将 Data Reader Op 的输出的数据进行解码,并输出

对于一些简单的数据格式,不需要解码,可以省略掉 Data Decoder,只使用 Data Reader 即可。

作为示例, Mini Dataloader 处理的数据格式虽然简单,但是依然实现了 DataReader 及 Data Decoder 两类 op,其中:

  • MiniReader 负责从文件中读取数据,并按逗号分隔字符串,将文本转为浮点数据后,设置到 Op 的输出中,输出形状为每行两列
  • MiniDecoder 负责将以上每行两列的输出进行分割,得到2个每行1列的输出 xy

test_mini_dataloader.py 中可以看到 Python 层次两者的使用:

    miniRecord = MiniReader(
        "./",
        batch_size=batch_size,
        data_part_num=2,
        part_name_suffix_length=3,
        random_shuffle=True,
        shuffle_after_epoch=True,
    )

    x, y = MiniDecoder(
            miniRecord, name="d1"
        )

以下,我们将介绍 C++ 层次如何实现 Data Reader 算子与 Data Decoder 算子。

Data Reader 算子

Data Reader 的类关系

我们需要实现一个继承自 DataReader 的类,该类包含了两个重要对象 loader_parser_,分别继承自 DatasetParser

  • loader_ 的工作是从文件系统中读取数据至缓冲区,Op 作者通过重写 Next 方法编写这部分的逻辑
  • parser_ 的工作是将缓冲区中的数据,设置到 Op 的输出中,Op 作者通过重写 Parser 方法编写这部分的逻辑

当 Data Reader Op 工作时,会调用 loader_ 中的相关方法打开文件系统中的文件,并调用 loader_Next 方法按照 Op 作者预定的逻辑从文件系统读取数据,然后,再调用 parser_Parser 方法,将数据设置到 Op 的输出中。

以下的伪代码展示了以上类关系和调用过程,实际代码比伪代码要复杂,并不是一一对应的关系:

class DataReader{
    void Read(user_op::KernelComputeContext* ctx) {
    // 运行到此处,已经启动了多线程加速数据处理
    loader->next();
    parser_->Parse();
  }
    Dataset* loader_;
    Parser*  parser_;
};

class MiniDataReader : DataReader{
    loader_ = new MiniDataSet;
    parser_ = new MiniParser;
};

class MiniDataset: Dataset {
  MiniDataset() {
    //在文件系统中找到数据集,并打开文件,初始化输入流
    //...
  }

  Next() {
    // 从输入流中读取数据的逻辑
  }
};

class MiniParser: Parser {
  void Parse(){
    // 将 DataSet 中的数据 设置到 Op 的输出中
  }
};
在 Data Reader Op 的 Kernel 中,会触发 DataReaderRead 方法,进而完成以上伪代码所展示的一连串操作。

以下我们针对 MiniReader 算子的真实代码进行解析。

Op 及 Kernel 注册

我们通过以下代码,注册了 MiniReader 的 Op:

REGISTER_CPU_ONLY_USER_OP("MiniReader")
    .Output("out")
    .Attr("data_dir", UserOpAttrType::kAtString)
    .Attr("data_part_num", UserOpAttrType::kAtInt32)
    .Attr<std::string>("part_name_prefix", UserOpAttrType::kAtString, "part-")
    .Attr<int32_t>("part_name_suffix_length", UserOpAttrType::kAtInt32, -1)
    .Attr("batch_size", UserOpAttrType::kAtInt32)
    .Attr<bool>("random_shuffle", UserOpAttrType::kAtBool, false)
    .Attr<bool>("shuffle_after_epoch", UserOpAttrType::kAtBool, false)
    .Attr<int64_t>("seed", UserOpAttrType::kAtInt64, -1)
    .Attr<int32_t>("shuffle_buffer_size", UserOpAttrType::kAtInt32, 1024)
    .SetTensorDescInferFn([](user_op::InferContext* ctx) -> Maybe<void> {
      //...
      *out_tensor->mut_shape() = Shape({local_batch_size, 2});
      *out_tensor->mut_data_type() = DataType::kDouble;
      //...
    })
    .SetGetSbpFn([](user_op::SbpContext* ctx) -> Maybe<void> {
      ctx->NewBuilder().Split(ctx->outputs(), 0).Build();
      //...
    });
可以看到,因为 Data Reader 是比较特殊的 Op,只有输出,没有输入(数据来自文件系统,而不是神经网络中的某个上游节点),因此我们只通过 Out 方法设置了输出,并在 SetTensorDescInferFn 设置了输出的性质为每行2列,数据类为 DataType::kDouble。同理,在设置 SetGetSbpFn 中设置 SBP Signature 时,只需要设置输出的 SBP 属性,我们将其设置为 Split(0)。

而设置的各种属性(data_dirdata_part_num 等),沿用了 OFRecord 数据集 中关于文件命名规范的要求,这使得我们在加载自定义文件格式时,可以复用 OneFlow 中已有的相关代码,让用户可以像 加载 OFRecord 数据集 那样,加载我们自定义格式的文件。

接着看这个 Op 的 Kernel 实现:

class MiniReaderKernel final : public user_op::OpKernel {
 public:
  //...

  std::shared_ptr<user_op::OpKernelState>
  CreateOpKernelState(user_op::KernelInitContext* ctx) override{
    std::shared_ptr<MiniReaderWrapper> reader(new MiniReaderWrapper(ctx));
    return reader;
  }

  void Compute(user_op::KernelComputeContext* ctx,
               user_op::OpKernelState* state) override {
    auto* reader = dynamic_cast<MiniReaderWrapper*>(state);
    reader->Read(ctx);
  }
  //...
};

REGISTER_USER_KERNEL("MiniReader")
    .SetCreateFn<MiniReaderKernel>()
    .SetIsMatchedHob((user_op::HobDeviceTag() == "cpu")
                     & (user_op::HobDataType("out", 0) == DataType::kDouble));
依据 自定义 Op 一文中的知识,我们知道 MiniReaderKernel::Compute 负责 Op 的运算逻辑。不过此处使用使用的 Compute 是包含2个参数的重载,我们有必要介绍下它的第二个参数 OpKernelState

当我们进行 Compute 时,有时除了从 ctx 获取的信息外,我们还需要维护一些其他的对象,这种对象不需要反复创建,但是其中的信息状态可能随着 Compute 多次调用而改变。为了应对这种需求,OneFlow 提供了2个参数的 Compute 的重载,为了使用它,我们必须同时重写 CreateOpKernelStateCreateOpKernelState 的作用是返回一个 user_op::OpKernelState 派生类对象,这个对象,将在 Compute 调用时,作为第二个参数传递。

为此,我们只需要将除 ctx 外想要维护的信息,封装为 user_op::OpKernelState 的派生类,并在 CreateOpKernelState 实例化并返回即可。

具体到我们实现的 Mini Reader 的 Kernel,我们先实现了一个继承自 user_op::OpKernelState 的类 MiniReaderWrapper, 它是对 MiniDataReader 的简单封装,之所以封装一层 MiniReaderWrapper 而不直接使用 MiniDataReader,仅仅是为了符合以上所述的 OneFlow 框架要求。

class MiniReaderWrapper final : public user_op::OpKernelState {
 public:
  explicit MiniReaderWrapper(user_op::KernelInitContext* ctx) : reader_(ctx) {}
  ~MiniReaderWrapper() = default;

  void Read(user_op::KernelComputeContext* ctx) { reader_.Read(ctx); }

 private:
  data::MiniDataReader reader_;
};

然后,重写 CreateOpKernelState,在其内部创建一个 MiniReaderWrapper 对象:

  std::shared_ptr<user_op::OpKernelState>
  CreateOpKernelState(user_op::KernelInitContext* ctx) override{
    std::shared_ptr<MiniReaderWrapper> reader(new MiniReaderWrapper(ctx));
    return reader;
  }

这样,在适当的时机,OneFlow 就会自动调用 CreateOpKernelState 创建对象,并将其作为第2个参数传递给 Compute。我们可以在 Compute 中拿到这个对象,并使用:

    auto* reader = dynamic_cast<MiniReaderWrapper*>(state);
    reader->Read(ctx);
可以看到,在 MiniReader 的 Kernel 中,我们仅仅是简单调用了 MiniReaderWrapper::Reader,这会触发上文伪代码中所提及的 DataReader::Read 流程。

MiniDataReader

上文伪代码中已经提及,在 MiniDataReader 内部,会实例化一个 MiniDataset 并赋值给 loader_ 指针。 以下是真实代码:

class MiniDataReader final : public DataReader<TensorBuffer> {
 public:
  MiniDataReader(user_op::KernelInitContext* ctx) : DataReader<TensorBuffer>(ctx) {
    loader_.reset(new MiniDataset(ctx));
    parser_.reset(new MiniParser());
    if (ctx->Attr<bool>("random_shuffle")) {
      loader_.reset(new RandomShuffleDataset<TensorBuffer>(ctx, std::move(loader_)));
    }
    int32_t batch_size = ctx->TensorDesc4ArgNameAndIndex("out", 0)->shape().elem_cnt();
    loader_.reset(new BatchDataset<TensorBuffer>(batch_size, std::move(loader_)));
    StartLoadThread();
  }
};
可以看到,除了我们自己继承自 DataSetMiniDataset 类之外,OneFlow 还内置了其他的 XXXDataSet,它们可以在已有的 DataSet 基础上增加额外功能,如 RandomShuffleDataset 用于 shuffle,BatchDataset 用于批量读取数据。 一切完成后,最后调用 StartLoadThread,顾名思义,启动加载线程,在 StartLoadThread 中,最终会触发重写的 MiniDataset::Next 方法。

以上 MiniDataReader 的构造,可以作为模板,没有特殊要求,在实现自定义的 DataLoader 过程中,不需要修改。

MiniDataset

对于 MiniDataSet,我们只需要关心它的构造函数以及重写的 Next 方法。

构造函数主要是通过 Attr 获取用户的配置,然后根据用户配置,初始化输入流。以下代码中的 JoinDirPath 内部,主要根据数据集文件名的约定(前缀、文件数目,文件名编号是否补齐等),获取所有的文件名称;而 InitInStream 是将数据集中的文件,都初始化为 OneFlow 封装的输入流(in_stream_ 成员),这在后续的 Next 方法中会使用。

  MiniDataset(user_op::KernelInitContext* ctx) {
    current_epoch_ = 0;
    shuffle_after_epoch_ = ctx->Attr<bool>("shuffle_after_epoch");

    //Join Dir Path
    JoinDirPath(ctx);

    // in stream
    InitInStream(ctx);
  }

从文件中加载的逻辑,写在 Next 虚函数中:

  LoadTargetPtrList Next() override {
    LoadTargetPtrList ret;
    LoadTargetPtr sample_ptr(new TensorBuffer());

    std::string sampleline;
    if (in_stream_->ReadLine(&sampleline) != 0) {
      ShuffleAfterEpoch();
      in_stream_->ReadLine(&sampleline);
    }

    auto numbers = CommaSplit(sampleline);
    sample_ptr->Resize(Shape({2}), DataType::kDouble);
    auto pNums = sample_ptr->mut_data<double>();
    pNums[0] = std::stod(numbers[0]);
    pNums[1] = std::stod(numbers[1]);
    ret.push_back(std::move(sample_ptr));

    return ret;
  }
在以上代码中,我们通过调用 in_stream_ReadLine 方法,将文件中的数据,读取至 string 对象 sampleline 中。然后通过 CommaSplit 等操作,将字符串按逗号分隔,并转为浮点数,放置到 TensorBuffer 对象中。

指的一提的是,in_stream_ 有2种方法从文件中读取数据,分别是:

int32_t PersistentInStream::ReadLine(std::string* l);
int32_t PersistentInStream::ReadFully(char* s, size_t n);
ReadLine 读取文件中的一行,至 l 对象;ReadFully 读取 n 个字节的数据,至 s 所指向的内存。均以0作为成功时的返回值。

MiniDataSet 完成从文件到内存缓冲区的工作,接着,我们将使用 MiniParser,将缓冲区中的内容,设置到 Op 的输出中。

MiniParser

MiniParser 继承自 Parser,我们只需要重写其中的 Parser 方法即可:

class MiniParser final : public Parser<TensorBuffer> {
 public:
  using LoadTargetPtr = std::shared_ptr<TensorBuffer>;
  using LoadTargetPtrList = std::vector<LoadTargetPtr>;

  void Parse(std::shared_ptr<LoadTargetPtrList> batch_data,
             user_op::KernelComputeContext* ctx) override {
    user_op::Tensor* out_tensor = ctx->Tensor4ArgNameAndIndex("out", 0);
    double* dptr = out_tensor->mut_dptr<double>();

    MultiThreadLoop(batch_data->size(), [&](size_t i) {
      TensorBuffer* buffer = batch_data->at(i).get();
      dptr[i*2]= *(buffer->data<double>());
      dptr[i*2+1]= *(buffer->data<double>()+1);
    });
  }
};
Parser 包含2个参数,其中 batch_data 其实是一个封装了的 vecotr,这个容器内的每个元素,就是之前 MiniDataSet 通过 Next 读取的数据。 参数 ctx 使得我们可以获取 Op 的信息,在这里,我们主要通过 ctx 获取输出,并获取指向输出缓冲区的指针 dptr

注意,我们将 batch_data 中的数据设置到 Op 的输出 dptr 的过程中,使用了宏 MultiThreadLoopMultiThreadLoop 可以让我们的循环逻辑在多线程中执行,它接受2个参数,第一个参数为循环的总次数;第二个参数是一个回调函数,原型为 void callback(size_t i),OneFlow 会创建多个线程,然后并发调用这个回调函数。回调函数的参数 i 表明了当前循环的序号,使得我们可以根据 i 来划分数据,完成自己的业务逻辑。

在以上的代码中,我们通过 batch_data->at(i).get() 获取了缓冲区的第 i 个的数据,然后将其设置到输出的内存区的第 i 行的位置,一共2列。

Data Decoder 算子

Data Decoder 算子,其实就是一种普通的算子,它接受 DataReader 的输出作为自己的输入,然后通过运算后,输出一个或者多个 Blob。

ofrecord_decoder_ops.cpp 可以看到针对 OFRecord 数据的各种 decoder。

我们的 Mini Dataloader 处理的数据比较简单,因此 MiniDecoder 所做的工作也很简单,仅仅是将 DataReader 所输出的每行2列的数据,拆分为2个每行1列的输出 xy

Mini Decoder 的 Op 注册为:

REGISTER_CPU_ONLY_USER_OP("mini_decoder")
    .Input("in")
    .Output("x")
    .Output("y")
    .SetTensorDescInferFn([](user_op::InferContext* ctx) -> Maybe<void> {
      user_op::TensorDesc* in_tensor = ctx->TensorDesc4ArgNameAndIndex("in", 0);
      user_op::TensorDesc* out_tensor_x = ctx->TensorDesc4ArgNameAndIndex("x", 0);
      user_op::TensorDesc* out_tensor_y = ctx->TensorDesc4ArgNameAndIndex("y", 0);
      // 设置输入输出 Blob 的属性
      // ...
    })
    .SetGetSbpFn([](user_op::SbpContext* ctx) -> Maybe<void> {
      ctx->NewBuilder()
          .Split(user_op::OpArg("in", 0), 0)
          .Split(user_op::OpArg("x", 0), 0)
          .Split(user_op::OpArg("y", 0), 0)
          .Build();
      //...
    });

Mini Decoder 的 Kernel 的实现:

class MiniDecoderKernel final : public user_op::OpKernel {
  //...
  void Compute(user_op::KernelComputeContext* ctx) const override {
    user_op::Tensor* in_blob = ctx->Tensor4ArgNameAndIndex("in", 0);
    user_op::Tensor* out_blob_x = ctx->Tensor4ArgNameAndIndex("x", 0);
    user_op::Tensor* out_blob_y = ctx->Tensor4ArgNameAndIndex("y", 0);

    int64_t record_num = in_blob->shape().At(0);

    const double* input = in_blob->dptr<double>();
    double* out_dptr_x = out_blob_x->mut_dptr<double>();
    double* out_dptr_y = out_blob_y->mut_dptr<double>();

    MultiThreadLoop(record_num, [&](size_t i){
      *(out_dptr_x + i) = *(input+i*2);
      *(out_dptr_y + i) = *(input+i*2 + 1);
    });

  }
  //...
};
可见,在 MiniDecoderKernel::Compute 中主要是获取到输入 in_blob, 然后在多线程循环 MultiThreadLoop 中,将输入的数据拆分到 out_dptr_xout_dptr_y 中,它们分别对应了输出 xy

自定义 DataLoader 的使用

自定义 Op 一文中所描述,要使用 C++ 层编写的 Op,还需要在 Python 层封装一个 Python Wrapper,才可以在 Python 层使用。这些工作放到了 test_mini_dataloader.py中:

def MiniDecoder(
    input_blob,
    name = None,
):
    if name is None:
        name = "Mini_Decoder_uniqueID"
    return (
        flow.user_op_builder(name)
        .Op("mini_decoder")
        .Input("in", [input_blob])
        .Output("x")
        .Output("y")
        .Build()
        .InferAndTryRun()
        .RemoteBlobList()
    )

def MiniReader(
    minidata_dir: str,
    batch_size: int = 1,
    data_part_num: int = 2,
    part_name_prefix: str = "part-",
    part_name_suffix_length: int = -1,
    random_shuffle: bool = False,
    shuffle_after_epoch: bool = False,
    shuffle_buffer_size: int = 1024,
    name = None,
):
    if name is None:
        name = "Mini_Reader_uniqueID"

    return (
        flow.user_op_builder(name)
        .Op("MiniReader")
        .Output("out")
        .Attr("data_dir", minidata_dir)
        .Attr("data_part_num", data_part_num)
        .Attr("batch_size", batch_size)
        .Attr("part_name_prefix", part_name_prefix)
        .Attr("random_shuffle", random_shuffle)
        .Attr("shuffle_after_epoch", shuffle_after_epoch)
        .Attr("part_name_suffix_length", part_name_suffix_length)
        .Attr("shuffle_buffer_size", shuffle_buffer_size)
        .Build()
        .InferAndTryRun()
        .RemoteBlobList()[0]
    )

test_mini_dataloader.py 中,我们使用自己实现的 MiniReaderMiniDecoder 加载并解码了数据集(part-000part-001)中的数据,完成了一次训练。

Mini Dataloader 的编译与测试

进入到本文对应的 data_loader 目录。 修改 Makefile 文件中的 ONEFLOW_ROOT 变量为 OneFlow 源码路径。 然后通过

make
可生成 miniloader.so 文件。

然后运行 test_mini_dataloader.py 脚本,可以使用 Mini Dataloader 加载数据并完成训练。

python test_mini_dataloader.py