You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

datasets.cc 70 kB

5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830
  1. /**
  2. * Copyright 2020 Huawei Technologies Co., Ltd
  3. *
  4. * Licensed under the Apache License, Version 2.0 (the "License");
  5. * you may not use this file except in compliance with the License.
  6. * You may obtain a copy of the License at
  7. *
  8. * http://www.apache.org/licenses/LICENSE-2.0
  9. *
  10. * Unless required by applicable law or agreed to in writing, software
  11. * distributed under the License is distributed on an "AS IS" BASIS,
  12. * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. * See the License for the specific language governing permissions and
  14. * limitations under the License.
  15. */
  16. #include <fstream>
  17. #include "minddata/dataset/include/datasets.h"
  18. #include "minddata/dataset/include/samplers.h"
  19. #include "minddata/dataset/include/transforms.h"
  20. #include "minddata/dataset/engine/dataset_iterator.h"
  21. // Source dataset headers (in alphabetical order)
  22. #include "minddata/dataset/engine/datasetops/source/album_op.h"
  23. #include "minddata/dataset/engine/datasetops/source/celeba_op.h"
  24. #include "minddata/dataset/engine/datasetops/source/cifar_op.h"
  25. #include "minddata/dataset/engine/datasetops/source/clue_op.h"
  26. #include "minddata/dataset/engine/datasetops/source/coco_op.h"
  27. #include "minddata/dataset/engine/datasetops/source/csv_op.h"
  28. #include "minddata/dataset/engine/datasetops/source/image_folder_op.h"
  29. #include "minddata/dataset/engine/datasetops/source/manifest_op.h"
  30. #include "minddata/dataset/engine/datasetops/source/mnist_op.h"
  31. #include "minddata/dataset/engine/datasetops/source/random_data_op.h"
  32. #include "minddata/dataset/engine/datasetops/source/text_file_op.h"
  33. #include "minddata/dataset/engine/datasetops/source/voc_op.h"
  34. // Dataset operator headers (in alphabetical order)
  35. #include "minddata/dataset/engine/datasetops/batch_op.h"
  36. #include "minddata/dataset/engine/datasetops/build_vocab_op.h"
  37. #include "minddata/dataset/engine/datasetops/concat_op.h"
  38. #include "minddata/dataset/engine/datasetops/map_op/map_op.h"
  39. #include "minddata/dataset/engine/datasetops/project_op.h"
  40. #include "minddata/dataset/engine/datasetops/rename_op.h"
  41. #include "minddata/dataset/engine/datasetops/repeat_op.h"
  42. #include "minddata/dataset/engine/datasetops/shuffle_op.h"
  43. #include "minddata/dataset/engine/datasetops/skip_op.h"
  44. #include "minddata/dataset/engine/datasetops/take_op.h"
  45. #include "minddata/dataset/engine/datasetops/zip_op.h"
  46. // Sampler headers (in alphabetical order)
  47. #include "minddata/dataset/engine/datasetops/source/sampler/sampler.h"
  48. #include "minddata/dataset/engine/datasetops/source/sampler/random_sampler.h"
  49. #include "minddata/dataset/engine/datasetops/source/sampler/sequential_sampler.h"
  50. #include "minddata/dataset/core/config_manager.h"
  51. #include "minddata/dataset/util/random.h"
  52. #include "minddata/dataset/util/path.h"
  53. namespace mindspore {
  54. namespace dataset {
  55. namespace api {
  56. #define RETURN_EMPTY_IF_ERROR(_s) \
  57. do { \
  58. Status __rc = (_s); \
  59. if (__rc.IsError()) { \
  60. MS_LOG(ERROR) << __rc; \
  61. return {}; \
  62. } \
  63. } while (false)
  64. // Function to create the iterator, which will build and launch the execution tree.
  65. std::shared_ptr<Iterator> Dataset::CreateIterator(std::vector<std::string> columns) {
  66. std::shared_ptr<Iterator> iter;
  67. try {
  68. auto ds = shared_from_this();
  69. // The specified columns will be selected from the dataset and passed down the pipeline
  70. // in the order specified, other columns will be discarded.
  71. if (!columns.empty()) {
  72. ds = ds->Project(columns);
  73. }
  74. iter = std::make_shared<Iterator>();
  75. Status rc = iter->BuildAndLaunchTree(ds);
  76. if (rc.IsError()) {
  77. MS_LOG(ERROR) << "CreateIterator failed." << rc;
  78. return nullptr;
  79. }
  80. return iter;
  81. } catch (const std::exception &err) {
  82. MS_LOG(ERROR) << "CreateIterator: Iterator exception caught: " << err.what();
  83. return nullptr;
  84. }
  85. return iter;
  86. }
  87. // Constructor
  88. Dataset::Dataset() {
  89. // Fetch some default value from config manager
  90. std::shared_ptr<ConfigManager> cfg = GlobalContext::config_manager();
  91. num_workers_ = cfg->num_parallel_workers();
  92. rows_per_buffer_ = cfg->rows_per_buffer();
  93. connector_que_size_ = cfg->op_connector_size();
  94. worker_connector_size_ = cfg->worker_connector_size();
  95. }
  96. /// \brief Function to create a SchemaObj
  97. /// \param[in] schema_file Path of schema file
  98. /// \return Shared pointer to the current schema
  99. std::shared_ptr<SchemaObj> Schema(const std::string &schema_file) {
  100. auto schema = std::make_shared<SchemaObj>(schema_file);
  101. return schema->init() ? schema : nullptr;
  102. }
  103. // FUNCTIONS TO CREATE DATASETS FOR LEAF-NODE DATASETS
  104. // (In alphabetical order)
  105. // Function to create a AlbumDataset.
  106. std::shared_ptr<AlbumDataset> Album(const std::string &dataset_dir, const std::string &data_schema,
  107. const std::vector<std::string> &column_names, bool decode,
  108. const std::shared_ptr<SamplerObj> &sampler) {
  109. auto ds = std::make_shared<AlbumDataset>(dataset_dir, data_schema, column_names, decode, sampler);
  110. return ds->ValidateParams() ? ds : nullptr;
  111. }
  112. // Function to create a CelebADataset.
  113. std::shared_ptr<CelebADataset> CelebA(const std::string &dataset_dir, const std::string &dataset_type,
  114. const std::shared_ptr<SamplerObj> &sampler, bool decode,
  115. const std::set<std::string> &extensions) {
  116. auto ds = std::make_shared<CelebADataset>(dataset_dir, dataset_type, sampler, decode, extensions);
  117. // Call derived class validation method.
  118. return ds->ValidateParams() ? ds : nullptr;
  119. }
  120. // Function to create a Cifar10Dataset.
  121. std::shared_ptr<Cifar10Dataset> Cifar10(const std::string &dataset_dir, const std::shared_ptr<SamplerObj> &sampler) {
  122. auto ds = std::make_shared<Cifar10Dataset>(dataset_dir, sampler);
  123. // Call derived class validation method.
  124. return ds->ValidateParams() ? ds : nullptr;
  125. }
  126. // Function to create a Cifar100Dataset.
  127. std::shared_ptr<Cifar100Dataset> Cifar100(const std::string &dataset_dir, const std::shared_ptr<SamplerObj> &sampler) {
  128. auto ds = std::make_shared<Cifar100Dataset>(dataset_dir, sampler);
  129. // Call derived class validation method.
  130. return ds->ValidateParams() ? ds : nullptr;
  131. }
  132. // Function to create a CLUEDataset.
  133. std::shared_ptr<CLUEDataset> CLUE(const std::vector<std::string> &clue_files, const std::string &task,
  134. const std::string &usage, int64_t num_samples, ShuffleMode shuffle,
  135. int32_t num_shards, int32_t shard_id) {
  136. auto ds = std::make_shared<CLUEDataset>(clue_files, task, usage, num_samples, shuffle, num_shards, shard_id);
  137. // Call derived class validation method.
  138. return ds->ValidateParams() ? ds : nullptr;
  139. }
  140. // Function to create a CocoDataset.
  141. std::shared_ptr<CocoDataset> Coco(const std::string &dataset_dir, const std::string &annotation_file,
  142. const std::string &task, const bool &decode,
  143. const std::shared_ptr<SamplerObj> &sampler) {
  144. auto ds = std::make_shared<CocoDataset>(dataset_dir, annotation_file, task, decode, sampler);
  145. // Call derived class validation method.
  146. return ds->ValidateParams() ? ds : nullptr;
  147. }
  148. // Function to create a CSVDataset.
  149. std::shared_ptr<CSVDataset> CSV(const std::vector<std::string> &dataset_files, char field_delim,
  150. const std::vector<std::shared_ptr<CsvBase>> &column_defaults,
  151. const std::vector<std::string> &column_names, int64_t num_samples, ShuffleMode shuffle,
  152. int32_t num_shards, int32_t shard_id) {
  153. auto ds = std::make_shared<CSVDataset>(dataset_files, field_delim, column_defaults, column_names, num_samples,
  154. shuffle, num_shards, shard_id);
  155. // Call derived class validation method.
  156. return ds->ValidateParams() ? ds : nullptr;
  157. }
  158. // Function to create a ImageFolderDataset.
  159. std::shared_ptr<ImageFolderDataset> ImageFolder(const std::string &dataset_dir, bool decode,
  160. const std::shared_ptr<SamplerObj> &sampler,
  161. const std::set<std::string> &extensions,
  162. const std::map<std::string, int32_t> &class_indexing) {
  163. // This arg exists in ImageFolderOp, but not externalized (in Python API). The default value is false.
  164. bool recursive = false;
  165. // Create logical representation of ImageFolderDataset.
  166. auto ds = std::make_shared<ImageFolderDataset>(dataset_dir, decode, sampler, recursive, extensions, class_indexing);
  167. // Call derived class validation method.
  168. return ds->ValidateParams() ? ds : nullptr;
  169. }
  170. // Function to create a ManifestDataset.
  171. std::shared_ptr<ManifestDataset> Manifest(std::string dataset_file, std::string usage,
  172. std::shared_ptr<SamplerObj> sampler,
  173. const std::map<std::string, int32_t> &class_indexing, bool decode) {
  174. auto ds = std::make_shared<ManifestDataset>(dataset_file, usage, sampler, class_indexing, decode);
  175. // Call derived class validation method.
  176. return ds->ValidateParams() ? ds : nullptr;
  177. }
  178. // Function to create a MnistDataset.
  179. std::shared_ptr<MnistDataset> Mnist(const std::string &dataset_dir, const std::shared_ptr<SamplerObj> &sampler) {
  180. auto ds = std::make_shared<MnistDataset>(dataset_dir, sampler);
  181. // Call derived class validation method.
  182. return ds->ValidateParams() ? ds : nullptr;
  183. }
  184. // Function to overload "+" operator to concat two datasets
  185. std::shared_ptr<ConcatDataset> operator+(const std::shared_ptr<Dataset> &datasets1,
  186. const std::shared_ptr<Dataset> &datasets2) {
  187. std::shared_ptr<ConcatDataset> ds = std::make_shared<ConcatDataset>(std::vector({datasets2, datasets1}));
  188. // Call derived class validation method.
  189. return ds->ValidateParams() ? ds : nullptr;
  190. }
  191. // Function to create a TextFileDataset.
  192. std::shared_ptr<TextFileDataset> TextFile(const std::vector<std::string> &dataset_files, int64_t num_samples,
  193. ShuffleMode shuffle, int32_t num_shards, int32_t shard_id) {
  194. auto ds = std::make_shared<TextFileDataset>(dataset_files, num_samples, shuffle, num_shards, shard_id);
  195. // Call derived class validation method.
  196. return ds->ValidateParams() ? ds : nullptr;
  197. }
  198. // Function to create a VOCDataset.
  199. std::shared_ptr<VOCDataset> VOC(const std::string &dataset_dir, const std::string &task, const std::string &mode,
  200. const std::map<std::string, int32_t> &class_indexing, bool decode,
  201. const std::shared_ptr<SamplerObj> &sampler) {
  202. auto ds = std::make_shared<VOCDataset>(dataset_dir, task, mode, class_indexing, decode, sampler);
  203. // Call derived class validation method.
  204. return ds->ValidateParams() ? ds : nullptr;
  205. }
  206. // Function to create a ZipDataset.
  207. std::shared_ptr<ZipDataset> Zip(const std::vector<std::shared_ptr<Dataset>> &datasets) {
  208. auto ds = std::make_shared<ZipDataset>(datasets);
  209. // Call derived class validation method.
  210. return ds->ValidateParams() ? ds : nullptr;
  211. }
  212. // FUNCTIONS TO CREATE DATASETS FOR DATASET OPS
  213. // (In alphabetical order)
  214. // Function to create a Batch dataset
  215. std::shared_ptr<BatchDataset> Dataset::Batch(int32_t batch_size, bool drop_remainder) {
  216. // Default values
  217. std::vector<std::string> cols_to_map = {};
  218. std::map<std::string, std::pair<TensorShape, std::shared_ptr<Tensor>>> pad_map;
  219. bool pad = false;
  220. auto ds = std::make_shared<BatchDataset>(batch_size, drop_remainder, pad, cols_to_map, pad_map);
  221. if (!ds->ValidateParams()) {
  222. return nullptr;
  223. }
  224. ds->children.push_back(shared_from_this());
  225. return ds;
  226. }
  227. // Function to create a Vocab from dataset
  228. std::shared_ptr<Vocab> Dataset::BuildVocab(const std::vector<std::string> &columns,
  229. const std::pair<int64_t, int64_t> &freq_range, int64_t top_k,
  230. const std::vector<std::string> &special_tokens, bool special_first) {
  231. auto vocab = std::make_shared<Vocab>();
  232. auto ds = std::make_shared<BuildVocabDataset>(vocab, columns, freq_range, top_k, special_tokens, special_first);
  233. if (!ds->ValidateParams()) {
  234. return nullptr;
  235. }
  236. ds->children.push_back(shared_from_this());
  237. // Run tree here to starting building vocab
  238. std::shared_ptr<Iterator> iter = ds->CreateIterator();
  239. if (iter == nullptr) {
  240. MS_LOG(ERROR) << "Fail to run iterator in BuildVocab.";
  241. return nullptr;
  242. }
  243. // Finish building vocab by triggering GetNextRow
  244. std::unordered_map<std::string, std::shared_ptr<Tensor>> row;
  245. iter->GetNextRow(&row);
  246. if (vocab == nullptr) {
  247. MS_LOG(ERROR) << "Fail to build vocab.";
  248. return nullptr;
  249. }
  250. return vocab;
  251. }
  252. // Function to create a Concat dataset
  253. std::shared_ptr<ConcatDataset> Dataset::Concat(const std::vector<std::shared_ptr<Dataset>> &datasets) {
  254. auto ds = std::make_shared<ConcatDataset>(datasets);
  255. ds->children.push_back(shared_from_this());
  256. return ds->ValidateParams() ? ds : nullptr;
  257. }
  258. // Function to create a Map dataset.
  259. std::shared_ptr<MapDataset> Dataset::Map(std::vector<std::shared_ptr<TensorOperation>> operations,
  260. std::vector<std::string> input_columns,
  261. std::vector<std::string> output_columns,
  262. const std::vector<std::string> &project_columns) {
  263. auto ds = std::make_shared<MapDataset>(operations, input_columns, output_columns, project_columns);
  264. if (!ds->ValidateParams()) {
  265. return nullptr;
  266. }
  267. ds->children.push_back(shared_from_this());
  268. return ds;
  269. }
  270. // Function to create a ProjectDataset.
  271. std::shared_ptr<ProjectDataset> Dataset::Project(const std::vector<std::string> &columns) {
  272. auto ds = std::make_shared<ProjectDataset>(columns);
  273. // Call derived class validation method.
  274. if (!ds->ValidateParams()) {
  275. return nullptr;
  276. }
  277. ds->children.push_back(shared_from_this());
  278. return ds;
  279. }
  280. // Function to create a RenameDataset.
  281. std::shared_ptr<RenameDataset> Dataset::Rename(const std::vector<std::string> &input_columns,
  282. const std::vector<std::string> &output_columns) {
  283. auto ds = std::make_shared<RenameDataset>(input_columns, output_columns);
  284. // Call derived class validation method.
  285. if (!ds->ValidateParams()) {
  286. return nullptr;
  287. }
  288. ds->children.push_back(shared_from_this());
  289. return ds;
  290. }
  291. // Function to create Repeat dataset.
  292. std::shared_ptr<Dataset> Dataset::Repeat(int32_t count) {
  293. // Workaround for repeat == 1, do not inject repeat.
  294. if (count == 1) {
  295. return shared_from_this();
  296. }
  297. auto ds = std::make_shared<RepeatDataset>(count);
  298. if (!ds->ValidateParams()) {
  299. return nullptr;
  300. }
  301. ds->children.push_back(shared_from_this());
  302. return ds;
  303. }
  304. // Function to create a ShuffleOp
  305. std::shared_ptr<ShuffleDataset> Dataset::Shuffle(int32_t buffer_size) {
  306. // Pass in reshuffle_each_epoch with true
  307. auto ds = std::make_shared<ShuffleDataset>(buffer_size, true);
  308. if (!ds->ValidateParams()) {
  309. return nullptr;
  310. }
  311. ds->children.push_back(shared_from_this());
  312. return ds;
  313. }
  314. // Function to create a SkipDataset.
  315. std::shared_ptr<SkipDataset> Dataset::Skip(int32_t count) {
  316. auto ds = std::make_shared<SkipDataset>(count);
  317. // Call derived class validation method.
  318. if (!ds->ValidateParams()) {
  319. return nullptr;
  320. }
  321. ds->children.push_back(shared_from_this());
  322. return ds;
  323. }
  324. // Function to create a TakeDataset.
  325. std::shared_ptr<Dataset> Dataset::Take(int32_t count) {
  326. // If count is greater than the number of element in dataset or equal to -1,
  327. // all the element in dataset will be taken
  328. if (count == -1) {
  329. return shared_from_this();
  330. }
  331. auto ds = std::make_shared<TakeDataset>(count);
  332. // Call derived class validation method.
  333. if (!ds->ValidateParams()) {
  334. return nullptr;
  335. }
  336. ds->children.push_back(shared_from_this());
  337. return ds;
  338. }
  339. // Function to create a Zip dataset
  340. std::shared_ptr<ZipDataset> Dataset::Zip(const std::vector<std::shared_ptr<Dataset>> &datasets) {
  341. // Default values
  342. auto ds = std::make_shared<ZipDataset>(datasets);
  343. ds->children.push_back(shared_from_this());
  344. return ds->ValidateParams() ? ds : nullptr;
  345. }
  346. SchemaObj::SchemaObj(const std::string &schema_file) : schema_file_(schema_file), num_rows_(0), dataset_type_("") {}
  347. // SchemaObj init function
  348. bool SchemaObj::init() {
  349. if (schema_file_ != "") {
  350. Path schema_file(schema_file_);
  351. if (!schema_file.Exists()) {
  352. MS_LOG(ERROR) << "The file " << schema_file << " does not exist or permission denied!";
  353. return false;
  354. }
  355. nlohmann::json js;
  356. try {
  357. std::ifstream in(schema_file_);
  358. in >> js;
  359. } catch (const std::exception &err) {
  360. MS_LOG(ERROR) << "Schema file failed to load";
  361. return false;
  362. }
  363. return from_json(js);
  364. }
  365. return true;
  366. }
  367. // Function to add a column to schema with a mstype de_type
  368. bool SchemaObj::add_column(std::string name, TypeId de_type, std::vector<int32_t> shape) {
  369. nlohmann::json new_column;
  370. new_column["name"] = name;
  371. // if de_type is mstype
  372. DataType data_type = dataset::MSTypeToDEType(de_type);
  373. new_column["type"] = data_type.ToString();
  374. if (shape.size() > 0) {
  375. new_column["shape"] = shape;
  376. new_column["rank"] = shape.size();
  377. } else {
  378. new_column["rank"] = 1;
  379. }
  380. columns_.push_back(new_column);
  381. return true;
  382. }
  383. // Function to add a column to schema with a string de_type
  384. bool SchemaObj::add_column(std::string name, std::string de_type, std::vector<int32_t> shape) {
  385. nlohmann::json new_column;
  386. new_column["name"] = name;
  387. DataType data_type(de_type);
  388. new_column["type"] = data_type.ToString();
  389. if (shape.size() > 0) {
  390. new_column["shape"] = shape;
  391. new_column["rank"] = shape.size();
  392. } else {
  393. new_column["rank"] = 1;
  394. }
  395. columns_.push_back(new_column);
  396. return true;
  397. }
  398. std::string SchemaObj::to_json() {
  399. nlohmann::json json_file;
  400. json_file["columns"] = columns_;
  401. if (dataset_type_ != "") {
  402. json_file["datasetType"] = dataset_type_;
  403. }
  404. if (num_rows_ > 0) {
  405. json_file["numRows"] = num_rows_;
  406. }
  407. return json_file.dump(2);
  408. }
  409. bool SchemaObj::parse_column(nlohmann::json columns) {
  410. std::string name, de_type;
  411. std::vector<int32_t> shape;
  412. columns_.clear();
  413. if (columns.type() == nlohmann::json::value_t::array) {
  414. // reference to python list
  415. for (auto column : columns) {
  416. auto key_name = column.find("name");
  417. if (key_name == column.end()) {
  418. MS_LOG(ERROR) << "Column's name is missing";
  419. return false;
  420. }
  421. name = *key_name;
  422. auto key_type = column.find("type");
  423. if (key_type == column.end()) {
  424. MS_LOG(ERROR) << "Column's type is missing";
  425. return false;
  426. }
  427. de_type = *key_type;
  428. shape.clear();
  429. auto key_shape = column.find("shape");
  430. if (key_shape != column.end()) {
  431. shape.insert(shape.end(), (*key_shape).begin(), (*key_shape).end());
  432. }
  433. if (!add_column(name, de_type, shape)) {
  434. return false;
  435. }
  436. }
  437. } else if (columns.type() == nlohmann::json::value_t::object) {
  438. for (const auto &it_child : columns.items()) {
  439. name = it_child.key();
  440. auto key_type = it_child.value().find("type");
  441. if (key_type == it_child.value().end()) {
  442. MS_LOG(ERROR) << "Column's type is missing";
  443. return false;
  444. }
  445. de_type = *key_type;
  446. shape.clear();
  447. auto key_shape = it_child.value().find("shape");
  448. if (key_shape != it_child.value().end()) {
  449. shape.insert(shape.end(), (*key_shape).begin(), (*key_shape).end());
  450. }
  451. if (!add_column(name, de_type, shape)) {
  452. return false;
  453. }
  454. }
  455. } else {
  456. MS_LOG(ERROR) << "columns must be dict or list, columns contain name, type, shape(optional).";
  457. return false;
  458. }
  459. return true;
  460. }
  461. bool SchemaObj::from_json(nlohmann::json json_obj) {
  462. for (const auto &it_child : json_obj.items()) {
  463. if (it_child.key() == "datasetType") {
  464. dataset_type_ = it_child.value();
  465. } else if (it_child.key() == "numRows") {
  466. num_rows_ = it_child.value();
  467. } else if (it_child.key() == "columns") {
  468. if (!parse_column(it_child.value())) {
  469. MS_LOG(ERROR) << "parse columns failed";
  470. return false;
  471. }
  472. } else {
  473. MS_LOG(ERROR) << "Unknown field " << it_child.key();
  474. return false;
  475. }
  476. }
  477. if (columns_.empty()) {
  478. MS_LOG(ERROR) << "Columns are missing.";
  479. return false;
  480. }
  481. if (num_rows_ <= 0) {
  482. MS_LOG(ERROR) << "numRows must be greater than 0";
  483. return false;
  484. }
  485. return true;
  486. }
  487. // OTHER FUNCTIONS
  488. // Helper function to create default RandomSampler.
  489. std::shared_ptr<SamplerObj> CreateDefaultSampler() {
  490. const int32_t num_samples = 0; // 0 means to sample all ids.
  491. bool replacement = false;
  492. return std::make_shared<RandomSamplerObj>(replacement, num_samples);
  493. }
  494. // Helper function to compute a default shuffle size
  495. Status ComputeShuffleSize(int64_t num_files, int64_t num_devices, int64_t num_rows, int64_t total_rows,
  496. int64_t *shuffle_size) {
  497. const int64_t average_files_multiplier = 4;
  498. const int64_t shuffle_max = 10000;
  499. int64_t avg_rows_per_file = 0;
  500. // Adjust the num rows per shard if sharding was given
  501. if (num_devices > 0) {
  502. if (num_rows % num_devices == 0) {
  503. num_rows = num_rows / num_devices;
  504. } else {
  505. num_rows = (num_rows / num_devices) + 1;
  506. }
  507. }
  508. // Cap based on total rows directive. Some ops do not have this and give value of 0.
  509. if (total_rows > 0) {
  510. num_rows = std::min(num_rows, total_rows);
  511. }
  512. // get the average per file
  513. CHECK_FAIL_RETURN_UNEXPECTED(num_files != 0, "The size of dataset_files must greater than 0.");
  514. avg_rows_per_file = num_rows / num_files;
  515. *shuffle_size = std::max(avg_rows_per_file * average_files_multiplier, shuffle_max);
  516. return Status::OK();
  517. }
  518. // Helper function to inject a shuffle operator over top of current operator being built
  519. Status AddShuffleOp(int64_t num_files, int64_t num_devices, int64_t num_rows, int64_t total_rows,
  520. int32_t connector_que_size, int32_t rows_per_buffer, std::shared_ptr<DatasetOp> *shuffle_op) {
  521. std::shared_ptr<ShuffleOp> new_shuffle_op = nullptr;
  522. int64_t shuffle_size = 0;
  523. RETURN_EMPTY_IF_ERROR(ComputeShuffleSize(num_files, num_devices, num_rows, total_rows, &shuffle_size));
  524. MS_LOG(INFO) << "Dataset::AddShuffleOp - num_rows: " << num_rows << ", shuffle_size: " << shuffle_size;
  525. // Add the shuffle op
  526. *shuffle_op = std::make_shared<ShuffleOp>(shuffle_size, GetSeed(), connector_que_size, true, rows_per_buffer);
  527. return Status::OK();
  528. }
  529. // Helper function to validate dataset directory parameter
  530. bool ValidateDatasetDirParam(const std::string &dataset_name, std::string dataset_dir) {
  531. if (dataset_dir.empty()) {
  532. MS_LOG(ERROR) << dataset_name << ": dataset_dir is not specified.";
  533. return false;
  534. }
  535. Path dir(dataset_dir);
  536. if (!dir.IsDirectory()) {
  537. MS_LOG(ERROR) << dataset_name << ": dataset_dir: [" << dataset_dir << "] is an invalid directory path.";
  538. return false;
  539. }
  540. if (access(dataset_dir.c_str(), R_OK) == -1) {
  541. MS_LOG(ERROR) << dataset_name << ": No access to specified dataset path: " << dataset_dir;
  542. return false;
  543. }
  544. return true;
  545. }
  546. // Helper function to validate dataset dataset files parameter
  547. bool ValidateDatasetFilesParam(const std::string &dataset_name, const std::vector<std::string> &dataset_files) {
  548. if (dataset_files.empty()) {
  549. MS_LOG(ERROR) << dataset_name << ": dataset_files is not specified.";
  550. return false;
  551. }
  552. for (auto f : dataset_files) {
  553. Path dataset_file(f);
  554. if (!dataset_file.Exists()) {
  555. MS_LOG(ERROR) << dataset_name << ": dataset file: [" << f << "] is invalid or does not exist.";
  556. return false;
  557. }
  558. }
  559. return true;
  560. }
  561. // Helper function to validate dataset num_shards and shard_id parameters
  562. bool ValidateDatasetShardParams(const std::string &dataset_name, int32_t num_shards, int32_t shard_id) {
  563. if (num_shards <= 0) {
  564. MS_LOG(ERROR) << dataset_name << ": Invalid num_shards: " << num_shards;
  565. return false;
  566. }
  567. if (shard_id < 0 || shard_id >= num_shards) {
  568. MS_LOG(ERROR) << dataset_name << ": Invalid input, shard_id: " << shard_id << ", num_shards: " << num_shards;
  569. return false;
  570. }
  571. return true;
  572. }
  573. /* ####################################### Derived Dataset classes ################################# */
  574. // DERIVED DATASET CLASSES LEAF-NODE DATASETS
  575. // (In alphabetical order)
  576. // Constructor for AlbumDataset
  577. AlbumDataset::AlbumDataset(const std::string &dataset_dir, const std::string &data_schema,
  578. const std::vector<std::string> &column_names, bool decode,
  579. const std::shared_ptr<SamplerObj> &sampler)
  580. : dataset_dir_(dataset_dir),
  581. schema_path_(data_schema),
  582. column_names_(column_names),
  583. decode_(decode),
  584. sampler_(sampler) {}
  585. bool AlbumDataset::ValidateParams() {
  586. if (!ValidateDatasetDirParam("AlbumDataset", dataset_dir_)) {
  587. return false;
  588. }
  589. if (!ValidateDatasetFilesParam("AlbumDataset", {schema_path_})) {
  590. return false;
  591. }
  592. return true;
  593. }
  594. // Function to build AlbumDataset
  595. std::vector<std::shared_ptr<DatasetOp>> AlbumDataset::Build() {
  596. // A vector containing shared pointer to the Dataset Ops that this object will create
  597. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  598. // If user does not specify Sampler, create a default sampler, i.e., RandomSampler.
  599. if (sampler_ == nullptr) {
  600. sampler_ = CreateDefaultSampler();
  601. }
  602. auto schema = std::make_unique<DataSchema>();
  603. RETURN_EMPTY_IF_ERROR(schema->LoadSchemaFile(schema_path_, column_names_));
  604. // Argument that is not exposed to user in the API.
  605. std::set<std::string> extensions = {};
  606. node_ops.push_back(std::make_shared<AlbumOp>(num_workers_, rows_per_buffer_, dataset_dir_, connector_que_size_,
  607. decode_, extensions, std::move(schema), std::move(sampler_->Build())));
  608. return node_ops;
  609. }
  610. // Constructor for CelebADataset
  611. CelebADataset::CelebADataset(const std::string &dataset_dir, const std::string &dataset_type,
  612. const std::shared_ptr<SamplerObj> &sampler, const bool &decode,
  613. const std::set<std::string> &extensions)
  614. : dataset_dir_(dataset_dir),
  615. dataset_type_(dataset_type),
  616. sampler_(sampler),
  617. decode_(decode),
  618. extensions_(extensions) {}
  619. bool CelebADataset::ValidateParams() {
  620. if (!ValidateDatasetDirParam("CelebADataset", dataset_dir_)) {
  621. return false;
  622. }
  623. std::set<std::string> dataset_type_list = {"all", "train", "valid", "test"};
  624. auto iter = dataset_type_list.find(dataset_type_);
  625. if (iter == dataset_type_list.end()) {
  626. MS_LOG(ERROR) << "dataset_type should be one of 'all', 'train', 'valid' or 'test'.";
  627. return false;
  628. }
  629. return true;
  630. }
  631. // Function to build CelebADataset
  632. std::vector<std::shared_ptr<DatasetOp>> CelebADataset::Build() {
  633. // A vector containing shared pointer to the Dataset Ops that this object will create
  634. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  635. // If user does not specify Sampler, create a default sampler based on the shuffle variable.
  636. if (sampler_ == nullptr) {
  637. sampler_ = CreateDefaultSampler();
  638. }
  639. std::unique_ptr<DataSchema> schema = std::make_unique<DataSchema>();
  640. RETURN_EMPTY_IF_ERROR(
  641. schema->AddColumn(ColDescriptor("image", DataType(DataType::DE_UINT8), TensorImpl::kFlexible, 1)));
  642. // label is like this:0 1 0 0 1......
  643. RETURN_EMPTY_IF_ERROR(
  644. schema->AddColumn(ColDescriptor("attr", DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 1)));
  645. node_ops.push_back(std::make_shared<CelebAOp>(num_workers_, rows_per_buffer_, dataset_dir_, connector_que_size_,
  646. decode_, dataset_type_, extensions_, std::move(schema),
  647. std::move(sampler_->Build())));
  648. return node_ops;
  649. }
  650. // Constructor for Cifar10Dataset
  651. Cifar10Dataset::Cifar10Dataset(const std::string &dataset_dir, std::shared_ptr<SamplerObj> sampler)
  652. : dataset_dir_(dataset_dir), sampler_(sampler) {}
  653. bool Cifar10Dataset::ValidateParams() { return ValidateDatasetDirParam("Cifar10Dataset", dataset_dir_); }
  654. // Function to build CifarOp for Cifar10
  655. std::vector<std::shared_ptr<DatasetOp>> Cifar10Dataset::Build() {
  656. // A vector containing shared pointer to the Dataset Ops that this object will create
  657. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  658. // If user does not specify Sampler, create a default sampler based on the shuffle variable.
  659. if (sampler_ == nullptr) {
  660. sampler_ = CreateDefaultSampler();
  661. }
  662. // Do internal Schema generation.
  663. auto schema = std::make_unique<DataSchema>();
  664. RETURN_EMPTY_IF_ERROR(schema->AddColumn(ColDescriptor("image", DataType(DataType::DE_UINT8), TensorImpl::kCv, 1)));
  665. TensorShape scalar = TensorShape::CreateScalar();
  666. RETURN_EMPTY_IF_ERROR(
  667. schema->AddColumn(ColDescriptor("label", DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 0, &scalar)));
  668. node_ops.push_back(std::make_shared<CifarOp>(CifarOp::CifarType::kCifar10, num_workers_, rows_per_buffer_,
  669. dataset_dir_, connector_que_size_, std::move(schema),
  670. std::move(sampler_->Build())));
  671. return node_ops;
  672. }
  673. // Constructor for Cifar100Dataset
  674. Cifar100Dataset::Cifar100Dataset(const std::string &dataset_dir, std::shared_ptr<SamplerObj> sampler)
  675. : dataset_dir_(dataset_dir), sampler_(sampler) {}
  676. bool Cifar100Dataset::ValidateParams() { return ValidateDatasetDirParam("Cifar100Dataset", dataset_dir_); }
  677. // Function to build CifarOp for Cifar100
  678. std::vector<std::shared_ptr<DatasetOp>> Cifar100Dataset::Build() {
  679. // A vector containing shared pointer to the Dataset Ops that this object will create
  680. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  681. // If user does not specify Sampler, create a default sampler based on the shuffle variable.
  682. if (sampler_ == nullptr) {
  683. sampler_ = CreateDefaultSampler();
  684. }
  685. // Do internal Schema generation.
  686. auto schema = std::make_unique<DataSchema>();
  687. RETURN_EMPTY_IF_ERROR(schema->AddColumn(ColDescriptor("image", DataType(DataType::DE_UINT8), TensorImpl::kCv, 1)));
  688. TensorShape scalar = TensorShape::CreateScalar();
  689. RETURN_EMPTY_IF_ERROR(
  690. schema->AddColumn(ColDescriptor("coarse_label", DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 0, &scalar)));
  691. RETURN_EMPTY_IF_ERROR(
  692. schema->AddColumn(ColDescriptor("fine_label", DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 0, &scalar)));
  693. node_ops.push_back(std::make_shared<CifarOp>(CifarOp::CifarType::kCifar100, num_workers_, rows_per_buffer_,
  694. dataset_dir_, connector_que_size_, std::move(schema),
  695. std::move(sampler_->Build())));
  696. return node_ops;
  697. }
  698. // Constructor for CLUEDataset
  699. CLUEDataset::CLUEDataset(const std::vector<std::string> clue_files, std::string task, std::string usage,
  700. int64_t num_samples, ShuffleMode shuffle, int32_t num_shards, int32_t shard_id)
  701. : dataset_files_(clue_files),
  702. task_(task),
  703. usage_(usage),
  704. num_samples_(num_samples),
  705. shuffle_(shuffle),
  706. num_shards_(num_shards),
  707. shard_id_(shard_id) {}
  708. bool CLUEDataset::ValidateParams() {
  709. if (!ValidateDatasetFilesParam("CLUEDataset", dataset_files_)) {
  710. return false;
  711. }
  712. std::vector<std::string> task_list = {"AFQMC", "TNEWS", "IFLYTEK", "CMNLI", "WSC", "CSL"};
  713. std::vector<std::string> usage_list = {"train", "test", "eval"};
  714. if (find(task_list.begin(), task_list.end(), task_) == task_list.end()) {
  715. MS_LOG(ERROR) << "task should be AFQMC, TNEWS, IFLYTEK, CMNLI, WSC or CSL.";
  716. return false;
  717. }
  718. if (find(usage_list.begin(), usage_list.end(), usage_) == usage_list.end()) {
  719. MS_LOG(ERROR) << "usage should be train, test or eval.";
  720. return false;
  721. }
  722. if (num_samples_ < 0) {
  723. MS_LOG(ERROR) << "CLUEDataset: Invalid number of samples: " << num_samples_;
  724. return false;
  725. }
  726. if (!ValidateDatasetShardParams("CLUEDataset", num_shards_, shard_id_)) {
  727. return false;
  728. }
  729. return true;
  730. }
  731. // Function to split string based on a character delimiter
  732. std::vector<std::string> CLUEDataset::split(const std::string &s, char delim) {
  733. std::vector<std::string> res;
  734. std::stringstream ss(s);
  735. std::string item;
  736. while (getline(ss, item, delim)) {
  737. res.push_back(item);
  738. }
  739. return res;
  740. }
  741. // Function to build CLUEDataset
  742. std::vector<std::shared_ptr<DatasetOp>> CLUEDataset::Build() {
  743. // A vector containing shared pointer to the Dataset Ops that this object will create
  744. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  745. std::map<std::string, std::string> key_map;
  746. if (task_ == "AFQMC") {
  747. if (usage_ == "train") {
  748. key_map["sentence1"] = "sentence1";
  749. key_map["sentence2"] = "sentence2";
  750. key_map["label"] = "label";
  751. } else if (usage_ == "test") {
  752. key_map["id"] = "id";
  753. key_map["sentence1"] = "sentence1";
  754. key_map["sentence2"] = "sentence2";
  755. } else if (usage_ == "eval") {
  756. key_map["sentence1"] = "sentence1";
  757. key_map["sentence2"] = "sentence2";
  758. key_map["label"] = "label";
  759. }
  760. } else if (task_ == "CMNLI") {
  761. if (usage_ == "train") {
  762. key_map["sentence1"] = "sentence1";
  763. key_map["sentence2"] = "sentence2";
  764. key_map["label"] = "label";
  765. } else if (usage_ == "test") {
  766. key_map["id"] = "id";
  767. key_map["sentence1"] = "sentence1";
  768. key_map["sentence2"] = "sentence2";
  769. } else if (usage_ == "eval") {
  770. key_map["sentence1"] = "sentence1";
  771. key_map["sentence2"] = "sentence2";
  772. key_map["label"] = "label";
  773. }
  774. } else if (task_ == "CSL") {
  775. if (usage_ == "train") {
  776. key_map["id"] = "id";
  777. key_map["abst"] = "abst";
  778. key_map["keyword"] = "keyword";
  779. key_map["label"] = "label";
  780. } else if (usage_ == "test") {
  781. key_map["id"] = "id";
  782. key_map["abst"] = "abst";
  783. key_map["keyword"] = "keyword";
  784. } else if (usage_ == "eval") {
  785. key_map["id"] = "id";
  786. key_map["abst"] = "abst";
  787. key_map["keyword"] = "keyword";
  788. key_map["label"] = "label";
  789. }
  790. } else if (task_ == "IFLYTEK") {
  791. if (usage_ == "train") {
  792. key_map["label"] = "label";
  793. key_map["label_des"] = "label_des";
  794. key_map["sentence"] = "sentence";
  795. } else if (usage_ == "test") {
  796. key_map["id"] = "id";
  797. key_map["sentence"] = "sentence";
  798. } else if (usage_ == "eval") {
  799. key_map["label"] = "label";
  800. key_map["label_des"] = "label_des";
  801. key_map["sentence"] = "sentence";
  802. }
  803. } else if (task_ == "TNEWS") {
  804. if (usage_ == "train") {
  805. key_map["label"] = "label";
  806. key_map["label_desc"] = "label_desc";
  807. key_map["sentence"] = "sentence";
  808. key_map["keywords"] = "keywords";
  809. } else if (usage_ == "test") {
  810. key_map["id"] = "id";
  811. key_map["sentence"] = "sentence";
  812. key_map["keywords"] = "keywords";
  813. } else if (usage_ == "eval") {
  814. key_map["label"] = "label";
  815. key_map["label_desc"] = "label_desc";
  816. key_map["sentence"] = "sentence";
  817. key_map["keywords"] = "keywords";
  818. }
  819. } else if (task_ == "WSC") {
  820. if (usage_ == "train") {
  821. key_map["span1_index"] = "target/span1_index";
  822. key_map["span2_index"] = "target/span2_index";
  823. key_map["span1_text"] = "target/span1_text";
  824. key_map["span2_text"] = "target/span2_text";
  825. key_map["idx"] = "idx";
  826. key_map["label"] = "label";
  827. key_map["text"] = "text";
  828. } else if (usage_ == "test") {
  829. key_map["span1_index"] = "target/span1_index";
  830. key_map["span2_index"] = "target/span2_index";
  831. key_map["span1_text"] = "target/span1_text";
  832. key_map["span2_text"] = "target/span2_text";
  833. key_map["idx"] = "idx";
  834. key_map["text"] = "text";
  835. } else if (usage_ == "eval") {
  836. key_map["span1_index"] = "target/span1_index";
  837. key_map["span2_index"] = "target/span2_index";
  838. key_map["span1_text"] = "target/span1_text";
  839. key_map["span2_text"] = "target/span2_text";
  840. key_map["idx"] = "idx";
  841. key_map["label"] = "label";
  842. key_map["text"] = "text";
  843. }
  844. }
  845. ColKeyMap ck_map;
  846. for (auto &p : key_map) {
  847. ck_map.insert({p.first, split(p.second, '/')});
  848. }
  849. bool shuffle_files = (shuffle_ == ShuffleMode::kGlobal || shuffle_ == ShuffleMode::kFiles);
  850. // Sort the dataset files in a lexicographical order
  851. std::vector<std::string> sorted_dataset_files = dataset_files_;
  852. std::sort(sorted_dataset_files.begin(), sorted_dataset_files.end());
  853. std::shared_ptr<ClueOp> clue_op =
  854. std::make_shared<ClueOp>(num_workers_, rows_per_buffer_, num_samples_, worker_connector_size_, ck_map,
  855. sorted_dataset_files, connector_que_size_, shuffle_files, num_shards_, shard_id_);
  856. RETURN_EMPTY_IF_ERROR(clue_op->Init());
  857. if (shuffle_ == ShuffleMode::kGlobal) {
  858. // Inject ShuffleOp
  859. std::shared_ptr<DatasetOp> shuffle_op = nullptr;
  860. int64_t num_rows = 0;
  861. // First, get the number of rows in the dataset
  862. RETURN_EMPTY_IF_ERROR(ClueOp::CountAllFileRows(sorted_dataset_files, &num_rows));
  863. // Add the shuffle op after this op
  864. RETURN_EMPTY_IF_ERROR(AddShuffleOp(sorted_dataset_files.size(), num_shards_, num_rows, 0, connector_que_size_,
  865. rows_per_buffer_, &shuffle_op));
  866. node_ops.push_back(shuffle_op);
  867. }
  868. node_ops.push_back(clue_op);
  869. return node_ops;
  870. }
  871. // Constructor for CocoDataset
  872. CocoDataset::CocoDataset(const std::string &dataset_dir, const std::string &annotation_file, const std::string &task,
  873. const bool &decode, const std::shared_ptr<SamplerObj> &sampler)
  874. : dataset_dir_(dataset_dir), annotation_file_(annotation_file), task_(task), decode_(decode), sampler_(sampler) {}
  875. bool CocoDataset::ValidateParams() {
  876. if (!ValidateDatasetDirParam("CocoDataset", dataset_dir_)) {
  877. return false;
  878. }
  879. Path annotation_file(annotation_file_);
  880. if (!annotation_file.Exists()) {
  881. MS_LOG(ERROR) << "annotation_file is invalid or not exist";
  882. return false;
  883. }
  884. std::set<std::string> task_list = {"Detection", "Stuff", "Panoptic", "Keypoint"};
  885. auto task_iter = task_list.find(task_);
  886. if (task_iter == task_list.end()) {
  887. MS_LOG(ERROR) << "Invalid task type";
  888. return false;
  889. }
  890. return true;
  891. }
  892. // Function to build CocoDataset
  893. std::vector<std::shared_ptr<DatasetOp>> CocoDataset::Build() {
  894. // A vector containing shared pointer to the Dataset Ops that this object will create
  895. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  896. // If user does not specify Sampler, create a default sampler based on the shuffle variable.
  897. if (sampler_ == nullptr) {
  898. sampler_ = CreateDefaultSampler();
  899. }
  900. CocoOp::TaskType task_type;
  901. if (task_ == "Detection") {
  902. task_type = CocoOp::TaskType::Detection;
  903. } else if (task_ == "Stuff") {
  904. task_type = CocoOp::TaskType::Stuff;
  905. } else if (task_ == "Keypoint") {
  906. task_type = CocoOp::TaskType::Keypoint;
  907. } else if (task_ == "Panoptic") {
  908. task_type = CocoOp::TaskType::Panoptic;
  909. }
  910. std::unique_ptr<DataSchema> schema = std::make_unique<DataSchema>();
  911. RETURN_EMPTY_IF_ERROR(
  912. schema->AddColumn(ColDescriptor(std::string("image"), DataType(DataType::DE_UINT8), TensorImpl::kFlexible, 1)));
  913. switch (task_type) {
  914. case CocoOp::TaskType::Detection:
  915. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  916. ColDescriptor(std::string("bbox"), DataType(DataType::DE_FLOAT32), TensorImpl::kFlexible, 1)));
  917. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  918. ColDescriptor(std::string("category_id"), DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 1)));
  919. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  920. ColDescriptor(std::string("iscrowd"), DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 1)));
  921. break;
  922. case CocoOp::TaskType::Stuff:
  923. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  924. ColDescriptor(std::string("segmentation"), DataType(DataType::DE_FLOAT32), TensorImpl::kFlexible, 1)));
  925. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  926. ColDescriptor(std::string("iscrowd"), DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 1)));
  927. break;
  928. case CocoOp::TaskType::Keypoint:
  929. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  930. ColDescriptor(std::string("keypoints"), DataType(DataType::DE_FLOAT32), TensorImpl::kFlexible, 1)));
  931. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  932. ColDescriptor(std::string("num_keypoints"), DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 1)));
  933. break;
  934. case CocoOp::TaskType::Panoptic:
  935. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  936. ColDescriptor(std::string("bbox"), DataType(DataType::DE_FLOAT32), TensorImpl::kFlexible, 1)));
  937. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  938. ColDescriptor(std::string("category_id"), DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 1)));
  939. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  940. ColDescriptor(std::string("iscrowd"), DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 1)));
  941. RETURN_EMPTY_IF_ERROR(
  942. schema->AddColumn(ColDescriptor(std::string("area"), DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 1)));
  943. break;
  944. default:
  945. MS_LOG(ERROR) << "CocoDataset::Build : Invalid task type";
  946. return {};
  947. }
  948. std::shared_ptr<CocoOp> op =
  949. std::make_shared<CocoOp>(task_type, dataset_dir_, annotation_file_, num_workers_, rows_per_buffer_,
  950. connector_que_size_, decode_, std::move(schema), std::move(sampler_->Build()));
  951. node_ops.push_back(op);
  952. return node_ops;
  953. }
  954. // Constructor for CSVDataset
  955. CSVDataset::CSVDataset(const std::vector<std::string> &csv_files, char field_delim,
  956. const std::vector<std::shared_ptr<CsvBase>> &column_defaults,
  957. const std::vector<std::string> &column_names, int64_t num_samples, ShuffleMode shuffle,
  958. int32_t num_shards, int32_t shard_id)
  959. : dataset_files_(csv_files),
  960. field_delim_(field_delim),
  961. column_defaults_(column_defaults),
  962. column_names_(column_names),
  963. num_samples_(num_samples),
  964. shuffle_(shuffle),
  965. num_shards_(num_shards),
  966. shard_id_(shard_id) {}
  967. bool CSVDataset::ValidateParams() {
  968. if (!ValidateDatasetFilesParam("CSVDataset", dataset_files_)) {
  969. return false;
  970. }
  971. if (field_delim_ == '"' || field_delim_ == '\r' || field_delim_ == '\n') {
  972. MS_LOG(ERROR) << "CSVDataset: The field delimiter should not be \", \\r, \\n";
  973. return false;
  974. }
  975. if (num_samples_ < -1) {
  976. MS_LOG(ERROR) << "CSVDataset: Invalid number of samples: " << num_samples_;
  977. return false;
  978. }
  979. if (!ValidateDatasetShardParams("CSVDataset", num_shards_, shard_id_)) {
  980. return false;
  981. }
  982. return true;
  983. }
  984. // Function to build CSVDataset
  985. std::vector<std::shared_ptr<DatasetOp>> CSVDataset::Build() {
  986. // A vector containing shared pointer to the Dataset Ops that this object will create
  987. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  988. bool shuffle_files = (shuffle_ == ShuffleMode::kGlobal || shuffle_ == ShuffleMode::kFiles);
  989. // Sort the dataset files in a lexicographical order
  990. std::vector<std::string> sorted_dataset_files = dataset_files_;
  991. std::sort(sorted_dataset_files.begin(), sorted_dataset_files.end());
  992. std::vector<std::shared_ptr<CsvOp::BaseRecord>> column_default_list;
  993. for (auto v : column_defaults_) {
  994. if (v->type == CsvType::INT) {
  995. column_default_list.push_back(
  996. std::make_shared<CsvOp::Record<int>>(CsvOp::INT, std::dynamic_pointer_cast<CsvRecord<int>>(v)->value));
  997. } else if (v->type == CsvType::FLOAT) {
  998. column_default_list.push_back(
  999. std::make_shared<CsvOp::Record<float>>(CsvOp::FLOAT, std::dynamic_pointer_cast<CsvRecord<float>>(v)->value));
  1000. } else if (v->type == CsvType::STRING) {
  1001. column_default_list.push_back(std::make_shared<CsvOp::Record<std::string>>(
  1002. CsvOp::STRING, std::dynamic_pointer_cast<CsvRecord<std::string>>(v)->value));
  1003. }
  1004. }
  1005. std::shared_ptr<CsvOp> csv_op = std::make_shared<CsvOp>(
  1006. sorted_dataset_files, field_delim_, column_default_list, column_names_, num_workers_, rows_per_buffer_,
  1007. num_samples_, worker_connector_size_, connector_que_size_, shuffle_files, num_shards_, shard_id_);
  1008. RETURN_EMPTY_IF_ERROR(csv_op->Init());
  1009. if (shuffle_ == ShuffleMode::kGlobal) {
  1010. // Inject ShuffleOp
  1011. std::shared_ptr<DatasetOp> shuffle_op = nullptr;
  1012. int64_t num_rows = 0;
  1013. // First, get the number of rows in the dataset
  1014. RETURN_EMPTY_IF_ERROR(CsvOp::CountAllFileRows(sorted_dataset_files, column_names_.empty(), &num_rows));
  1015. // Add the shuffle op after this op
  1016. RETURN_EMPTY_IF_ERROR(AddShuffleOp(sorted_dataset_files.size(), num_shards_, num_rows, 0, connector_que_size_,
  1017. rows_per_buffer_, &shuffle_op));
  1018. node_ops.push_back(shuffle_op);
  1019. }
  1020. node_ops.push_back(csv_op);
  1021. return node_ops;
  1022. }
  1023. ImageFolderDataset::ImageFolderDataset(std::string dataset_dir, bool decode, std::shared_ptr<SamplerObj> sampler,
  1024. bool recursive, std::set<std::string> extensions,
  1025. std::map<std::string, int32_t> class_indexing)
  1026. : dataset_dir_(dataset_dir),
  1027. decode_(decode),
  1028. sampler_(sampler),
  1029. recursive_(recursive),
  1030. class_indexing_(class_indexing),
  1031. exts_(extensions) {}
  1032. bool ImageFolderDataset::ValidateParams() { return ValidateDatasetDirParam("ImageFolderDataset", dataset_dir_); }
  1033. std::vector<std::shared_ptr<DatasetOp>> ImageFolderDataset::Build() {
  1034. // A vector containing shared pointer to the Dataset Ops that this object will create
  1035. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1036. // If user does not specify Sampler, create a default sampler, i.e., RandomSampler.
  1037. if (sampler_ == nullptr) {
  1038. sampler_ = CreateDefaultSampler();
  1039. }
  1040. // Do internal Schema generation.
  1041. // This arg is exist in ImageFolderOp, but not externalized (in Python API).
  1042. std::unique_ptr<DataSchema> schema = std::make_unique<DataSchema>();
  1043. TensorShape scalar = TensorShape::CreateScalar();
  1044. RETURN_EMPTY_IF_ERROR(
  1045. schema->AddColumn(ColDescriptor("image", DataType(DataType::DE_UINT8), TensorImpl::kFlexible, 1)));
  1046. RETURN_EMPTY_IF_ERROR(
  1047. schema->AddColumn(ColDescriptor("label", DataType(DataType::DE_INT32), TensorImpl::kFlexible, 0, &scalar)));
  1048. node_ops.push_back(std::make_shared<ImageFolderOp>(num_workers_, rows_per_buffer_, dataset_dir_, connector_que_size_,
  1049. recursive_, decode_, exts_, class_indexing_, std::move(schema),
  1050. std::move(sampler_->Build())));
  1051. return node_ops;
  1052. }
  1053. ManifestDataset::ManifestDataset(std::string dataset_file, std::string usage, std::shared_ptr<SamplerObj> sampler,
  1054. const std::map<std::string, int32_t> &class_indexing, bool decode)
  1055. : dataset_file_(dataset_file), usage_(usage), decode_(decode), class_index_(class_indexing), sampler_(sampler) {}
  1056. bool ManifestDataset::ValidateParams() {
  1057. Path manifest_file(dataset_file_);
  1058. if (!manifest_file.Exists()) {
  1059. MS_LOG(ERROR) << "dataset file: [" << dataset_file_ << "] is invalid or not exist";
  1060. return false;
  1061. }
  1062. std::vector<std::string> usage_list = {"train", "eval", "inference"};
  1063. if (find(usage_list.begin(), usage_list.end(), usage_) == usage_list.end()) {
  1064. MS_LOG(ERROR) << "usage should be train, eval or inference.";
  1065. return false;
  1066. }
  1067. return true;
  1068. }
  1069. std::vector<std::shared_ptr<DatasetOp>> ManifestDataset::Build() {
  1070. // A vector containing shared pointer to the Dataset Ops that this object will create
  1071. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1072. // If user does not specify Sampler, create a default sampler based on the shuffle variable.
  1073. if (sampler_ == nullptr) {
  1074. sampler_ = CreateDefaultSampler();
  1075. }
  1076. // Do internal Schema generation.
  1077. auto schema = std::make_unique<DataSchema>();
  1078. RETURN_EMPTY_IF_ERROR(schema->AddColumn(ColDescriptor("image", DataType(DataType::DE_UINT8), TensorImpl::kCv, 1)));
  1079. TensorShape scalar = TensorShape::CreateScalar();
  1080. RETURN_EMPTY_IF_ERROR(
  1081. schema->AddColumn(ColDescriptor("label", DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 0, &scalar)));
  1082. std::shared_ptr<ManifestOp> manifest_op;
  1083. manifest_op =
  1084. std::make_shared<ManifestOp>(num_workers_, rows_per_buffer_, dataset_file_, connector_que_size_, decode_,
  1085. class_index_, std::move(schema), std::move(sampler_->Build()), usage_);
  1086. node_ops.push_back(manifest_op);
  1087. return node_ops;
  1088. }
  1089. MnistDataset::MnistDataset(std::string dataset_dir, std::shared_ptr<SamplerObj> sampler)
  1090. : dataset_dir_(dataset_dir), sampler_(sampler) {}
  1091. bool MnistDataset::ValidateParams() { return ValidateDatasetDirParam("MnistDataset", dataset_dir_); }
  1092. std::vector<std::shared_ptr<DatasetOp>> MnistDataset::Build() {
  1093. // A vector containing shared pointer to the Dataset Ops that this object will create
  1094. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1095. // If user does not specify Sampler, create a default sampler, i.e., RandomSampler.
  1096. if (sampler_ == nullptr) {
  1097. sampler_ = CreateDefaultSampler();
  1098. }
  1099. // Do internal Schema generation.
  1100. auto schema = std::make_unique<DataSchema>();
  1101. RETURN_EMPTY_IF_ERROR(schema->AddColumn(ColDescriptor("image", DataType(DataType::DE_UINT8), TensorImpl::kCv, 1)));
  1102. TensorShape scalar = TensorShape::CreateScalar();
  1103. RETURN_EMPTY_IF_ERROR(
  1104. schema->AddColumn(ColDescriptor("label", DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 0, &scalar)));
  1105. node_ops.push_back(std::make_shared<MnistOp>(num_workers_, rows_per_buffer_, dataset_dir_, connector_que_size_,
  1106. std::move(schema), std::move(sampler_->Build())));
  1107. return node_ops;
  1108. }
  1109. // ValideParams for RandomDataset
  1110. bool RandomDataset::ValidateParams() {
  1111. if (total_rows_ < 0) {
  1112. MS_LOG(ERROR) << "RandomDataset: total_rows must be greater than 0, now get " << total_rows_;
  1113. return false;
  1114. }
  1115. return true;
  1116. }
  1117. int32_t RandomDataset::GenRandomInt(int32_t min, int32_t max) {
  1118. std::uniform_int_distribution<int32_t> uniDist(min, max);
  1119. return uniDist(rand_gen_);
  1120. }
  1121. // Build for RandomDataset
  1122. std::vector<std::shared_ptr<DatasetOp>> RandomDataset::Build() {
  1123. // A vector containing shared pointer to the Dataset Ops that this object will create
  1124. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1125. rand_gen_.seed(GetSeed()); // seed the random generator
  1126. // If total rows was not given, then randomly pick a number
  1127. std::shared_ptr<SchemaObj> schema_obj;
  1128. if (!schema_path_.empty()) schema_obj = std::make_shared<SchemaObj>(schema_path_);
  1129. if (schema_obj != nullptr && total_rows_ == 0) {
  1130. total_rows_ = schema_obj->get_num_rows();
  1131. }
  1132. // If user does not specify Sampler, create a default sampler based on the shuffle variable.
  1133. if (sampler_ == nullptr) {
  1134. sampler_ = CreateDefaultSampler();
  1135. }
  1136. std::string schema_json_string, schema_file_path;
  1137. if (schema_ != nullptr) {
  1138. schema_->set_dataset_type("Random");
  1139. if (total_rows_ != 0) {
  1140. schema_->set_num_rows(total_rows_);
  1141. }
  1142. schema_json_string = schema_->to_json();
  1143. } else {
  1144. schema_file_path = schema_path_;
  1145. }
  1146. std::unique_ptr<DataSchema> data_schema;
  1147. std::vector<std::string> columns_to_load;
  1148. if (!schema_file_path.empty() || !schema_json_string.empty()) {
  1149. data_schema = std::make_unique<DataSchema>();
  1150. if (!schema_file_path.empty()) {
  1151. data_schema->LoadSchemaFile(schema_file_path, columns_to_load);
  1152. } else if (!schema_json_string.empty()) {
  1153. data_schema->LoadSchemaString(schema_json_string, columns_to_load);
  1154. }
  1155. }
  1156. std::shared_ptr<RandomDataOp> op;
  1157. op = std::make_shared<RandomDataOp>(num_workers_, connector_que_size_, rows_per_buffer_, total_rows_,
  1158. std::move(data_schema), std::move(sampler_->Build()));
  1159. node_ops.push_back(op);
  1160. return node_ops;
  1161. }
  1162. // Constructor for TextFileDataset
  1163. TextFileDataset::TextFileDataset(std::vector<std::string> dataset_files, int32_t num_samples, ShuffleMode shuffle,
  1164. int32_t num_shards, int32_t shard_id)
  1165. : dataset_files_(dataset_files),
  1166. num_samples_(num_samples),
  1167. shuffle_(shuffle),
  1168. num_shards_(num_shards),
  1169. shard_id_(shard_id) {}
  1170. bool TextFileDataset::ValidateParams() {
  1171. if (!ValidateDatasetFilesParam("TextFileDataset", dataset_files_)) {
  1172. return false;
  1173. }
  1174. if (num_samples_ < 0) {
  1175. MS_LOG(ERROR) << "TextFileDataset: Invalid number of samples: " << num_samples_;
  1176. return false;
  1177. }
  1178. if (!ValidateDatasetShardParams("TextFileDataset", num_shards_, shard_id_)) {
  1179. return false;
  1180. }
  1181. return true;
  1182. }
  1183. // Function to build TextFileDataset
  1184. std::vector<std::shared_ptr<DatasetOp>> TextFileDataset::Build() {
  1185. // A vector containing shared pointer to the Dataset Ops that this object will create
  1186. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1187. bool shuffle_files = (shuffle_ == ShuffleMode::kGlobal || shuffle_ == ShuffleMode::kFiles);
  1188. // Sort the dataset files in a lexicographical order
  1189. std::vector<std::string> sorted_dataset_files = dataset_files_;
  1190. std::sort(sorted_dataset_files.begin(), sorted_dataset_files.end());
  1191. // Do internal Schema generation.
  1192. auto schema = std::make_unique<DataSchema>();
  1193. RETURN_EMPTY_IF_ERROR(
  1194. schema->AddColumn(ColDescriptor("text", DataType(DataType::DE_UINT8), TensorImpl::kFlexible, 1)));
  1195. // Create and initalize TextFileOp
  1196. std::shared_ptr<TextFileOp> text_file_op = std::make_shared<TextFileOp>(
  1197. num_workers_, rows_per_buffer_, num_samples_, worker_connector_size_, std::move(schema), sorted_dataset_files,
  1198. connector_que_size_, shuffle_files, num_shards_, shard_id_, std::move(nullptr));
  1199. RETURN_EMPTY_IF_ERROR(text_file_op->Init());
  1200. if (shuffle_ == ShuffleMode::kGlobal) {
  1201. // Inject ShuffleOp
  1202. std::shared_ptr<DatasetOp> shuffle_op = nullptr;
  1203. int64_t num_rows = 0;
  1204. // First, get the number of rows in the dataset
  1205. RETURN_EMPTY_IF_ERROR(TextFileOp::CountAllFileRows(sorted_dataset_files, &num_rows));
  1206. // Add the shuffle op after this op
  1207. RETURN_EMPTY_IF_ERROR(AddShuffleOp(sorted_dataset_files.size(), num_shards_, num_rows, 0, connector_que_size_,
  1208. rows_per_buffer_, &shuffle_op));
  1209. node_ops.push_back(shuffle_op);
  1210. }
  1211. // Add TextFileOp
  1212. node_ops.push_back(text_file_op);
  1213. return node_ops;
  1214. }
  1215. // Constructor for VOCDataset
  1216. VOCDataset::VOCDataset(const std::string &dataset_dir, const std::string &task, const std::string &mode,
  1217. const std::map<std::string, int32_t> &class_indexing, bool decode,
  1218. std::shared_ptr<SamplerObj> sampler)
  1219. : dataset_dir_(dataset_dir),
  1220. task_(task),
  1221. mode_(mode),
  1222. class_index_(class_indexing),
  1223. decode_(decode),
  1224. sampler_(sampler) {}
  1225. bool VOCDataset::ValidateParams() {
  1226. Path dir(dataset_dir_);
  1227. if (!dir.IsDirectory()) {
  1228. MS_LOG(ERROR) << "Invalid dataset path or no dataset path is specified.";
  1229. return false;
  1230. }
  1231. if (task_ == "Segmentation") {
  1232. if (!class_index_.empty()) {
  1233. MS_LOG(ERROR) << "class_indexing is invalid in Segmentation task.";
  1234. return false;
  1235. }
  1236. Path imagesets_file = dir / "ImageSets" / "Segmentation" / mode_ + ".txt";
  1237. if (!imagesets_file.Exists()) {
  1238. MS_LOG(ERROR) << "Invalid mode: " << mode_ << ", file \"" << imagesets_file << "\" is not exists!";
  1239. return false;
  1240. }
  1241. } else if (task_ == "Detection") {
  1242. Path imagesets_file = dir / "ImageSets" / "Main" / mode_ + ".txt";
  1243. if (!imagesets_file.Exists()) {
  1244. MS_LOG(ERROR) << "Invalid mode: " << mode_ << ", file \"" << imagesets_file << "\" is not exists!";
  1245. return false;
  1246. }
  1247. } else {
  1248. MS_LOG(ERROR) << "Invalid task: " << task_;
  1249. return false;
  1250. }
  1251. return true;
  1252. }
  1253. // Function to build VOCDataset
  1254. std::vector<std::shared_ptr<DatasetOp>> VOCDataset::Build() {
  1255. // A vector containing shared pointer to the Dataset Ops that this object will create
  1256. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1257. // If user does not specify Sampler, create a default sampler based on the shuffle variable.
  1258. if (sampler_ == nullptr) {
  1259. sampler_ = CreateDefaultSampler();
  1260. }
  1261. auto schema = std::make_unique<DataSchema>();
  1262. VOCOp::TaskType task_type_;
  1263. if (task_ == "Segmentation") {
  1264. task_type_ = VOCOp::TaskType::Segmentation;
  1265. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  1266. ColDescriptor(std::string(kColumnImage), DataType(DataType::DE_UINT8), TensorImpl::kFlexible, 1)));
  1267. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  1268. ColDescriptor(std::string(kColumnTarget), DataType(DataType::DE_UINT8), TensorImpl::kFlexible, 1)));
  1269. } else if (task_ == "Detection") {
  1270. task_type_ = VOCOp::TaskType::Detection;
  1271. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  1272. ColDescriptor(std::string(kColumnImage), DataType(DataType::DE_UINT8), TensorImpl::kFlexible, 1)));
  1273. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  1274. ColDescriptor(std::string(kColumnBbox), DataType(DataType::DE_FLOAT32), TensorImpl::kFlexible, 1)));
  1275. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  1276. ColDescriptor(std::string(kColumnLabel), DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 1)));
  1277. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  1278. ColDescriptor(std::string(kColumnDifficult), DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 1)));
  1279. RETURN_EMPTY_IF_ERROR(schema->AddColumn(
  1280. ColDescriptor(std::string(kColumnTruncate), DataType(DataType::DE_UINT32), TensorImpl::kFlexible, 1)));
  1281. }
  1282. std::shared_ptr<VOCOp> voc_op;
  1283. voc_op = std::make_shared<VOCOp>(task_type_, mode_, dataset_dir_, class_index_, num_workers_, rows_per_buffer_,
  1284. connector_que_size_, decode_, std::move(schema), std::move(sampler_->Build()));
  1285. node_ops.push_back(voc_op);
  1286. return node_ops;
  1287. }
  1288. // DERIVED DATASET CLASSES LEAF-NODE DATASETS
  1289. // (In alphabetical order)
  1290. BatchDataset::BatchDataset(int32_t batch_size, bool drop_remainder, bool pad, std::vector<std::string> cols_to_map,
  1291. std::map<std::string, std::pair<TensorShape, std::shared_ptr<Tensor>>> pad_map)
  1292. : batch_size_(batch_size),
  1293. drop_remainder_(drop_remainder),
  1294. pad_(pad),
  1295. cols_to_map_(cols_to_map),
  1296. pad_map_(pad_map) {}
  1297. std::vector<std::shared_ptr<DatasetOp>> BatchDataset::Build() {
  1298. // A vector containing shared pointer to the Dataset Ops that this object will create
  1299. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1300. #ifdef ENABLE_PYTHON
  1301. py::function noop;
  1302. node_ops.push_back(std::make_shared<BatchOp>(batch_size_, drop_remainder_, pad_, connector_que_size_, num_workers_,
  1303. cols_to_map_, noop, noop, pad_map_));
  1304. #else
  1305. node_ops.push_back(std::make_shared<BatchOp>(batch_size_, drop_remainder_, pad_, connector_que_size_, num_workers_,
  1306. cols_to_map_, pad_map_));
  1307. #endif
  1308. return node_ops;
  1309. }
  1310. bool BatchDataset::ValidateParams() {
  1311. if (batch_size_ <= 0) {
  1312. MS_LOG(ERROR) << "Batch: batch_size should be positive integer, but got: " << batch_size_;
  1313. return false;
  1314. }
  1315. return true;
  1316. }
  1317. BuildVocabDataset::BuildVocabDataset(std::shared_ptr<Vocab> vocab, const std::vector<std::string> &columns,
  1318. const std::pair<int64_t, int64_t> &freq_range, int64_t top_k,
  1319. const std::vector<std::string> &special_tokens, bool special_first)
  1320. : vocab_(vocab),
  1321. columns_(columns),
  1322. freq_range_(freq_range),
  1323. top_k_(top_k),
  1324. special_tokens_(special_tokens),
  1325. special_first_(special_first) {}
  1326. // Function to build BuildVocabDataset
  1327. std::vector<std::shared_ptr<DatasetOp>> BuildVocabDataset::Build() {
  1328. // A vector containing shared pointer to the Dataset Ops that this object will create
  1329. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1330. std::shared_ptr<BuildVocabOp> build_vocab_op;
  1331. build_vocab_op = std::make_shared<BuildVocabOp>(vocab_, columns_, freq_range_, top_k_, special_tokens_,
  1332. special_first_, num_workers_, connector_que_size_);
  1333. node_ops.push_back(build_vocab_op);
  1334. return node_ops;
  1335. }
  1336. bool BuildVocabDataset::ValidateParams() {
  1337. if (vocab_ == nullptr) {
  1338. MS_LOG(ERROR) << "BuildVocab: vocab is null.";
  1339. return false;
  1340. }
  1341. if (top_k_ < 0) {
  1342. MS_LOG(ERROR) << "BuildVocab: top_k shoule be positive, but got: " << top_k_;
  1343. return false;
  1344. }
  1345. if (freq_range_.first < 0 || freq_range_.second > kDeMaxFreq || freq_range_.first > freq_range_.second) {
  1346. MS_LOG(ERROR) << "BuildVocab: requency_range [a,b] should be 0 <= a <= b (a,b are inclusive), "
  1347. << "but got [" << freq_range_.first << ", " << freq_range_.second << "]";
  1348. return false;
  1349. }
  1350. return true;
  1351. }
  1352. // Function to build ConcatOp
  1353. ConcatDataset::ConcatDataset(const std::vector<std::shared_ptr<Dataset>> &datasets) : datasets_(datasets) {
  1354. this->children = datasets_;
  1355. }
  1356. bool ConcatDataset::ValidateParams() {
  1357. if (datasets_.empty()) {
  1358. MS_LOG(ERROR) << "Concat: concatenated datasets are not specified.";
  1359. return false;
  1360. }
  1361. if (find(datasets_.begin(), datasets_.end(), nullptr) != datasets_.end()) {
  1362. MS_LOG(ERROR) << "Concat: concatenated dataset should not be null.";
  1363. return false;
  1364. }
  1365. return true;
  1366. }
  1367. std::vector<std::shared_ptr<DatasetOp>> ConcatDataset::Build() {
  1368. // A vector containing shared pointer to the Dataset Ops that this object will create
  1369. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1370. node_ops.push_back(std::make_shared<ConcatOp>(connector_que_size_));
  1371. return node_ops;
  1372. }
  1373. MapDataset::MapDataset(std::vector<std::shared_ptr<TensorOperation>> operations, std::vector<std::string> input_columns,
  1374. std::vector<std::string> output_columns, const std::vector<std::string> &project_columns)
  1375. : operations_(operations),
  1376. input_columns_(input_columns),
  1377. output_columns_(output_columns),
  1378. project_columns_(project_columns) {}
  1379. std::vector<std::shared_ptr<DatasetOp>> MapDataset::Build() {
  1380. // A vector containing shared pointer to the Dataset Ops that this object will create
  1381. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1382. std::vector<std::shared_ptr<TensorOp>> tensor_ops;
  1383. // Build tensorOp from tensorOperation vector
  1384. // This is to ensure each iterator hold its own copy of the tensorOp objects.
  1385. (void)std::transform(
  1386. operations_.begin(), operations_.end(), std::back_inserter(tensor_ops),
  1387. [](std::shared_ptr<TensorOperation> operation) -> std::shared_ptr<TensorOp> { return operation->Build(); });
  1388. // This parameter will be removed with next rebase
  1389. std::vector<std::string> col_orders;
  1390. auto map_op = std::make_shared<MapOp>(input_columns_, output_columns_, tensor_ops, num_workers_, connector_que_size_);
  1391. if (!project_columns_.empty()) {
  1392. auto project_op = std::make_shared<ProjectOp>(project_columns_);
  1393. node_ops.push_back(project_op);
  1394. }
  1395. node_ops.push_back(map_op);
  1396. return node_ops;
  1397. }
  1398. bool MapDataset::ValidateParams() {
  1399. if (operations_.empty()) {
  1400. MS_LOG(ERROR) << "Map: No operation is specified.";
  1401. return false;
  1402. }
  1403. return true;
  1404. }
  1405. // Function to build ProjectOp
  1406. ProjectDataset::ProjectDataset(const std::vector<std::string> &columns) : columns_(columns) {}
  1407. bool ProjectDataset::ValidateParams() {
  1408. if (columns_.empty()) {
  1409. MS_LOG(ERROR) << "No columns are specified.";
  1410. return false;
  1411. }
  1412. return true;
  1413. }
  1414. std::vector<std::shared_ptr<DatasetOp>> ProjectDataset::Build() {
  1415. // A vector containing shared pointer to the Dataset Ops that this object will create
  1416. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1417. node_ops.push_back(std::make_shared<ProjectOp>(columns_));
  1418. return node_ops;
  1419. }
  1420. // Function to build RenameOp
  1421. RenameDataset::RenameDataset(const std::vector<std::string> &input_columns,
  1422. const std::vector<std::string> &output_columns)
  1423. : input_columns_(input_columns), output_columns_(output_columns) {}
  1424. bool RenameDataset::ValidateParams() {
  1425. if (input_columns_.empty() || output_columns_.empty()) {
  1426. MS_LOG(ERROR) << "input and output columns must be specified";
  1427. return false;
  1428. }
  1429. if (input_columns_.size() != output_columns_.size()) {
  1430. MS_LOG(ERROR) << "input and output columns must be the same size";
  1431. return false;
  1432. }
  1433. for (uint32_t i = 0; i < input_columns_.size(); ++i) {
  1434. if (input_columns_[i].empty()) {
  1435. MS_LOG(ERROR) << "input_columns: column name should not be empty.";
  1436. return false;
  1437. }
  1438. if (output_columns_[i].empty()) {
  1439. MS_LOG(ERROR) << "output_columns: column name should not be empty.";
  1440. return false;
  1441. }
  1442. }
  1443. return true;
  1444. }
  1445. std::vector<std::shared_ptr<DatasetOp>> RenameDataset::Build() {
  1446. // A vector containing shared pointer to the Dataset Ops that this object will create
  1447. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1448. node_ops.push_back(std::make_shared<RenameOp>(input_columns_, output_columns_, connector_que_size_));
  1449. return node_ops;
  1450. }
  1451. RepeatDataset::RepeatDataset(int32_t count) : repeat_count_(count) {}
  1452. std::vector<std::shared_ptr<DatasetOp>> RepeatDataset::Build() {
  1453. // A vector containing shared pointer to the Dataset Ops that this object will create
  1454. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1455. node_ops.push_back(std::make_shared<RepeatOp>(repeat_count_));
  1456. return node_ops;
  1457. }
  1458. bool RepeatDataset::ValidateParams() {
  1459. if (repeat_count_ <= 0 && repeat_count_ != -1) {
  1460. MS_LOG(ERROR) << "Repeat: repeat_count should be either -1 or positive integer, repeat_count_: " << repeat_count_;
  1461. return false;
  1462. }
  1463. return true;
  1464. }
  1465. // Constructor for ShuffleDataset
  1466. ShuffleDataset::ShuffleDataset(int32_t shuffle_size, bool reset_every_epoch)
  1467. : shuffle_size_(shuffle_size), shuffle_seed_(GetSeed()), reset_every_epoch_(reset_every_epoch) {}
  1468. // Function to build the ShuffleOp
  1469. std::vector<std::shared_ptr<DatasetOp>> ShuffleDataset::Build() {
  1470. // A vector containing shared pointer to the Dataset Ops that this object will create
  1471. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1472. node_ops.push_back(std::make_shared<ShuffleOp>(shuffle_size_, shuffle_seed_, connector_que_size_, reset_every_epoch_,
  1473. rows_per_buffer_));
  1474. return node_ops;
  1475. }
  1476. // Function to validate the parameters for ShuffleDataset
  1477. bool ShuffleDataset::ValidateParams() {
  1478. if (shuffle_size_ <= 1) {
  1479. MS_LOG(ERROR) << "ShuffleDataset: Invalid input, shuffle_size: " << shuffle_size_;
  1480. return false;
  1481. }
  1482. return true;
  1483. }
  1484. // Constructor for SkipDataset
  1485. SkipDataset::SkipDataset(int32_t count) : skip_count_(count) {}
  1486. // Function to build the SkipOp
  1487. std::vector<std::shared_ptr<DatasetOp>> SkipDataset::Build() {
  1488. // A vector containing shared pointer to the Dataset Ops that this object will create
  1489. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1490. node_ops.push_back(std::make_shared<SkipOp>(skip_count_, connector_que_size_));
  1491. return node_ops;
  1492. }
  1493. // Function to validate the parameters for SkipDataset
  1494. bool SkipDataset::ValidateParams() {
  1495. if (skip_count_ <= -1) {
  1496. MS_LOG(ERROR) << "Skip: skip_count should not be negative, skip_count: " << skip_count_;
  1497. return false;
  1498. }
  1499. return true;
  1500. }
  1501. // Constructor for TakeDataset
  1502. TakeDataset::TakeDataset(int32_t count) : take_count_(count) {}
  1503. // Function to build the TakeOp
  1504. std::vector<std::shared_ptr<DatasetOp>> TakeDataset::Build() {
  1505. // A vector containing shared pointer to the Dataset Ops that this object will create
  1506. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1507. node_ops.push_back(std::make_shared<TakeOp>(take_count_, connector_que_size_));
  1508. return node_ops;
  1509. }
  1510. // Function to validate the parameters for TakeDataset
  1511. bool TakeDataset::ValidateParams() {
  1512. if (take_count_ <= 0 && take_count_ != -1) {
  1513. MS_LOG(ERROR) << "Take: take_count should be either -1 or positive integer, take_count: " << take_count_;
  1514. return false;
  1515. }
  1516. return true;
  1517. }
  1518. // Function to build ZipOp
  1519. ZipDataset::ZipDataset(const std::vector<std::shared_ptr<Dataset>> &datasets) : datasets_(datasets) {
  1520. for (auto dataset : datasets_) {
  1521. this->children.push_back(dataset);
  1522. }
  1523. }
  1524. bool ZipDataset::ValidateParams() {
  1525. if (datasets_.empty()) {
  1526. MS_LOG(ERROR) << "Zip: dataset to zip are not specified.";
  1527. return false;
  1528. }
  1529. return true;
  1530. }
  1531. std::vector<std::shared_ptr<DatasetOp>> ZipDataset::Build() {
  1532. // A vector containing shared pointer to the Dataset Ops that this object will create
  1533. std::vector<std::shared_ptr<DatasetOp>> node_ops;
  1534. node_ops.push_back(std::make_shared<ZipOp>(rows_per_buffer_, connector_que_size_));
  1535. return node_ops;
  1536. }
  1537. } // namespace api
  1538. } // namespace dataset
  1539. } // namespace mindspore