Method COCODataReader

oneflow/user/data/coco_data_reader.cpp:29–56

namespace data {

COCODataReader::COCODataReader(user_op::KernelInitContext* ctx) : DataReader<COCOImage>(ctx) {
  batch_size_ = ctx->TensorDesc4ArgNameAndIndex("image", 0)->shape().elem_cnt();
  if (auto* pool = TensorBufferPool::TryGet()) { pool->IncreasePoolSizeByBase(batch_size_); }

  std::shared_ptr<const COCOMeta> meta(new COCOMeta(
      ctx->Attr<int64_t>("session_id"), ctx->Attr<std::string>("annotation_file"),
      ctx->Attr<std::string>("image_dir"), ctx->Attr<bool>("remove_images_without_annotations")));
  std::unique_ptr<RandomAccessDataset<COCOImage>> coco_dataset_ptr(new COCODataset(ctx, meta));

  size_t world_size = 1;
  int64_t rank = 0;
  CHECK_JUST(InitDataSourceDistributedInfo(ctx, world_size, rank));
  loader_.reset(new DistributedTrainingDataset<COCOImage>(
      world_size, rank, ctx->Attr<bool>("stride_partition"), ctx->Attr<bool>("shuffle_after_epoch"),
      ctx->Attr<int64_t>("random_seed"), std::move(coco_dataset_ptr)));

  if (ctx->Attr<bool>("group_by_ratio")) {
    auto GetGroupId = [](const COCOImage& sample) {
      return static_cast<int64_t>(sample.height / sample.width);
    };
    loader_.reset(new GroupBatchDataset<COCOImage>(batch_size_, GetGroupId, std::move(loader_)));
  } else {
    loader_.reset(new BatchDataset<COCOImage>(batch_size_, std::move(loader_)));
  }

  parser_.reset(new COCOParser(meta));
  StartLoadThread();
}

COCODataReader::~COCODataReader() {
  if (auto* pool = TensorBufferPool::TryGet()) { pool->DecreasePoolSizeByBase(batch_size_); }