dataset.py 44.3 KB
Newer Older
D
dongdaxiang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13
#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
T
tianshuo78520a 已提交
14
"""This is definition of dataset class, which is high performance IO."""
D
dongdaxiang 已提交
15 16 17 18

from paddle.fluid.proto import data_feed_pb2
from google.protobuf import text_format
from . import core
19
from ..utils import deprecated
20

D
dongdaxiang 已提交
21
__all__ = ['DatasetFactory', 'InMemoryDataset', 'QueueDataset']
D
dongdaxiang 已提交
22 23


24
class DatasetFactory:
25 26
    """
    DatasetFactory is a factory which create dataset by its name,
H
hutuxian 已提交
27
    you can create "QueueDataset" or "InMemoryDataset", or "FileInstantDataset",
28 29 30
    the default is "QueueDataset".

    Example:
31 32 33 34 35
        .. code-block:: python

          import paddle.fluid as fluid
          dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")

36
    """
D
dongdaxiang 已提交
37

D
dongdaxiang 已提交
38
    def __init__(self):
39
        """Init."""
D
dongdaxiang 已提交
40 41
        pass

42
    def create_dataset(self, datafeed_class="QueueDataset"):
43
        """
H
hutuxian 已提交
44
        Create "QueueDataset" or "InMemoryDataset", or "FileInstantDataset",
45
        the default is "QueueDataset".
D
dongdaxiang 已提交
46

47 48 49 50
        Args:
            datafeed_class(str): datafeed class name, QueueDataset or InMemoryDataset.
                                 Default is QueueDataset.

D
dongdaxiang 已提交
51
        Examples:
52 53 54 55 56
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset()

57
        """
D
dongdaxiang 已提交
58 59
        try:
            dataset = globals()[datafeed_class]()
60
            return dataset
D
dongdaxiang 已提交
61
        except:
62 63 64
            raise ValueError(
                "datafeed class %s does not exist" % datafeed_class
            )
D
dongdaxiang 已提交
65 66


67
class DatasetBase:
68
    """Base dataset class."""
D
dongdaxiang 已提交
69

D
dongdaxiang 已提交
70
    def __init__(self):
71
        """Init."""
D
dongdaxiang 已提交
72 73 74 75
        # define class name here
        # to decide whether we need create in memory instance
        self.proto_desc = data_feed_pb2.DataFeedDesc()
        self.proto_desc.pipe_command = "cat"
X
xujiaqi01 已提交
76
        self.dataset = core.Dataset("MultiSlotDataset")
77
        self.thread_num = 1
J
jiaqi 已提交
78
        self.filelist = []
79
        self.use_ps_gpu = False
80
        self.psgpu = None
D
dongdaxiang 已提交
81 82 83 84 85 86

    def set_pipe_command(self, pipe_command):
        """
        Set pipe command of current dataset
        A pipe command is a UNIX pipeline command that can be used only

87 88 89 90 91 92
        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset()
              dataset.set_pipe_command("python my_script.py")
93 94

        Args:
95
            pipe_command(str): pipe command
96

D
dongdaxiang 已提交
97 98 99
        """
        self.proto_desc.pipe_command = pipe_command

T
Thunderbrook 已提交
100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116
    def set_so_parser_name(self, so_parser_name):
        """
        Set so parser name of current dataset

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset()
              dataset.set_so_parser_name("./abc.so")

        Args:
            pipe_command(str): pipe command

        """
        self.proto_desc.so_parser_name = so_parser_name

117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133
    def set_rank_offset(self, rank_offset):
        """
        Set rank_offset for merge_pv. It set the message of Pv.

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset()
              dataset.set_rank_offset("rank_offset")

        Args:
            rank_offset(str): rank_offset's name

        """
        self.proto_desc.rank_offset = rank_offset

134 135 136 137
    def set_fea_eval(self, record_candidate_size, fea_eval=True):
        """
        set fea eval mode for slots shuffle to debug the importance level of
        slots(features), fea_eval need to be set True for slots shuffle.
138

139
        Args:
140
            record_candidate_size(int): size of instances candidate to shuffle
141
                                        one slot
T
tianshuo78520a 已提交
142
            fea_eval(bool): whether enable fea eval mode to enable slots shuffle.
143
                            default is True.
144

145 146 147 148 149 150 151 152 153 154 155 156 157 158
        Examples:
            .. code-block:: python

            import paddle.fluid as fluid
            dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
            dataset.set_fea_eval(1000000, True)

        """
        if fea_eval:
            self.dataset.set_fea_eval(fea_eval, record_candidate_size)
        self.fea_eval = fea_eval

    def slots_shuffle(self, slots):
        """
159 160
        Slots Shuffle
        Slots Shuffle is a shuffle method in slots level, which is usually used
161
        in sparse feature with large scale of instances. To compare the metric, i.e.
162
        auc while doing slots shuffle on one or several slots with baseline to
163
        evaluate the importance level of slots(features).
164

165 166 167 168 169 170 171 172 173 174 175 176 177 178
        Args:
            slots(list[string]): the set of slots(string) to do slots shuffle.

        Examples:
            import paddle.fluid as fluid
            dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
            dataset.set_merge_by_lineid()
            #suppose there is a slot 0
            dataset.slots_shuffle(['0'])
        """
        if self.fea_eval:
            slots_set = set(slots)
            self.dataset.slots_shuffle(slots_set)

D
dongdaxiang 已提交
179 180 181 182
    def set_batch_size(self, batch_size):
        """
        Set batch size. Will be effective during training

183 184 185 186 187 188
        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset()
              dataset.set_batch_size(128)
D
dongdaxiang 已提交
189 190

        Args:
191
            batch_size(int): batch size
D
dongdaxiang 已提交
192 193 194 195

        """
        self.proto_desc.batch_size = batch_size

196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211
    def set_pv_batch_size(self, pv_batch_size):
        """
        Set pv batch size. It will be effective during enable_pv_merge

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset()
              dataset.set_pv_batch(128)
        Args:
            pv_batch_size(int): pv batch size

        """
        self.proto_desc.pv_batch_size = pv_batch_size

212
    def set_thread(self, thread_num):
213 214 215
        """
        Set thread num, it is the num of readers.

216 217 218 219 220 221
        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset()
               dataset.set_thread(12)
222 223

        Args:
224
            thread_num(int): thread num
225
        """
226
        self.dataset.set_thread_num(thread_num)
227
        self.thread_num = thread_num
228 229

    def set_filelist(self, filelist):
230 231 232
        """
        Set file list in current worker.

233 234 235 236 237 238
        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset()
              dataset.set_filelist(['a.txt', 'b.txt'])
239 240

        Args:
241
            filelist(list): file list
242
        """
243
        self.dataset.set_filelist(filelist)
J
jiaqi 已提交
244
        self.filelist = filelist
245

246 247 248
    def set_input_type(self, input_type):
        self.proto_desc.input_type = input_type

D
dongdaxiang 已提交
249
    def set_use_var(self, var_list):
250 251 252
        """
        Set Variables which you will use.

253 254 255 256 257 258
        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset()
              dataset.set_use_var([data, label])
259 260

        Args:
261
            var_list(list): variable list
262
        """
263
        multi_slot = self.proto_desc.multi_slot_desc
D
dongdaxiang 已提交
264
        for var in var_list:
265
            slot_var = multi_slot.slots.add()
D
dongdaxiang 已提交
266 267 268 269
            slot_var.is_used = True
            slot_var.name = var.name
            if var.lod_level == 0:
                slot_var.is_dense = True
270
                slot_var.shape.extend(var.shape)
271
            if var.dtype == core.VarDesc.VarType.FP32:
D
dongdaxiang 已提交
272
                slot_var.type = "float"
273
            elif var.dtype == core.VarDesc.VarType.INT64:
D
dongdaxiang 已提交
274
                slot_var.type = "uint64"
B
Baibaifan 已提交
275 276
            elif var.dtype == core.VarDesc.VarType.INT32:
                slot_var.type = "uint32"
D
dongdaxiang 已提交
277 278
            else:
                raise ValueError(
B
Baibaifan 已提交
279
                    "Currently, fluid.dataset only supports dtype=float32, dtype=int32 and dtype=int64"
D
dongdaxiang 已提交
280 281
                )

282
    def set_hdfs_config(self, fs_name, fs_ugi):
283 284 285
        """
        Set hdfs config: fs name ad ugi

286 287 288 289 290 291
        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset()
              dataset.set_hdfs_config("my_fs_name", "my_fs_ugi")
292 293

        Args:
294 295
            fs_name(str): fs name
            fs_ugi(str): fs ugi
296
        """
297 298
        self.dataset.set_hdfs_config(fs_name, fs_ugi)

299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314
    def set_download_cmd(self, download_cmd):
        """
        Set customized download cmd: download_cmd

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset()
              dataset.set_download_cmd("./read_from_afs")

        Args:
            download_cmd(str): customized download command
        """
        self.dataset.set_download_cmd(download_cmd)

315
    def _prepare_to_run(self):
316 317 318 319
        """
        Set data_feed_desc before load or shuffle,
        user no need to call this function.
        """
J
jiaqi 已提交
320 321 322
        if self.thread_num > len(self.filelist):
            self.thread_num = len(self.filelist)
        self.dataset.set_thread_num(self.thread_num)
323
        self.dataset.set_data_feed_desc(self.desc())
J
jiaqi 已提交
324 325
        self.dataset.create_readers()

T
Thunderbrook 已提交
326
    def _set_use_ps_gpu(self, psgpu):
327 328 329 330 331 332
        """
        set use_ps_gpu flag

        Args:
            use_ps_gpu: bool
        """
T
Thunderbrook 已提交
333
        self.use_ps_gpu = True
334 335
        # if not defined heterps with paddle, users will not use psgpu
        if not core._is_compiled_with_heterps():
T
Thunderbrook 已提交
336
            self.use_ps_gpu = False
337
        elif self.use_ps_gpu:
T
Thunderbrook 已提交
338
            self.psgpu = psgpu
339

J
jiaqi 已提交
340 341
    def _finish_to_run(self):
        self.dataset.destroy_readers()
342

D
dongdaxiang 已提交
343 344 345 346
    def desc(self):
        """
        Returns a protobuf message for this DataFeedDesc

347 348 349 350 351 352
        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset()
              print(dataset.desc())
D
dongdaxiang 已提交
353 354 355 356 357 358

        Returns:
            A string message
        """
        return text_format.MessageToString(self.proto_desc)

359 360 361 362 363 364
    def _dynamic_adjust_before_train(self, thread_num):
        pass

    def _dynamic_adjust_after_train(self):
        pass

D
dongdaxiang 已提交
365 366

class InMemoryDataset(DatasetBase):
367 368
    """
    InMemoryDataset, it will load data into memory
D
dongdaxiang 已提交
369 370
    and shuffle data before training.
    This class should be created by DatasetFactory
371 372

    Example:
373
        dataset = paddle.fluid.DatasetFactory().create_dataset("InMemoryDataset")
374
    """
D
dongdaxiang 已提交
375

376
    @deprecated(since="2.0.0", update_to="paddle.distributed.InMemoryDataset")
D
dongdaxiang 已提交
377
    def __init__(self):
378
        """Init."""
379
        super().__init__()
380
        self.proto_desc.name = "MultiSlotInMemoryDataFeed"
381
        self.fleet_send_batch_size = None
382
        self.is_user_set_queue_num = False
J
jiaqi 已提交
383
        self.queue_num = None
384 385
        self.parse_ins_id = False
        self.parse_content = False
386 387 388
        self.parse_logkey = False
        self.merge_by_sid = True
        self.enable_pv_merge = False
389
        self.merge_by_lineid = False
390
        self.fleet_send_sleep_seconds = None
391
        self.trainer_num = -1
L
lxsbupt 已提交
392
        self.pass_id = 0
J
jiaqi 已提交
393

394 395 396 397
    @deprecated(
        since="2.0.0",
        update_to="paddle.distributed.InMemoryDataset._set_feed_type",
    )
398 399 400 401 402
    def set_feed_type(self, data_feed_type):
        """
        Set data_feed_desc
        """
        self.proto_desc.name = data_feed_type
403
        if self.proto_desc.name == "SlotRecordInMemoryDataFeed":
Y
yaoxuefeng 已提交
404
            self.dataset = core.Dataset("SlotRecordDataset")
405

406 407 408 409
    @deprecated(
        since="2.0.0",
        update_to="paddle.distributed.InMemoryDataset._prepare_to_run",
    )
J
jiaqi 已提交
410 411 412 413 414
    def _prepare_to_run(self):
        """
        Set data_feed_desc before load or shuffle,
        user no need to call this function.
        """
415
        if self.thread_num <= 0:
416
            self.thread_num = 1
J
jiaqi 已提交
417 418 419 420
        self.dataset.set_thread_num(self.thread_num)
        if self.queue_num is None:
            self.queue_num = self.thread_num
        self.dataset.set_queue_num(self.queue_num)
421 422
        self.dataset.set_parse_ins_id(self.parse_ins_id)
        self.dataset.set_parse_content(self.parse_content)
423 424 425
        self.dataset.set_parse_logkey(self.parse_logkey)
        self.dataset.set_merge_by_sid(self.merge_by_sid)
        self.dataset.set_enable_pv_merge(self.enable_pv_merge)
J
jiaqi 已提交
426 427 428 429
        self.dataset.set_data_feed_desc(self.desc())
        self.dataset.create_channel()
        self.dataset.create_readers()

430 431
    @deprecated(
        since="2.0.0",
432 433
        update_to="paddle.distributed.InMemoryDataset._dynamic_adjust_before_train",
    )
434 435
    def _dynamic_adjust_before_train(self, thread_num):
        if not self.is_user_set_queue_num:
436 437 438 439
            if self.use_ps_gpu:
                self.dataset.dynamic_adjust_channel_num(thread_num, True)
            else:
                self.dataset.dynamic_adjust_channel_num(thread_num, False)
440 441
        self.dataset.dynamic_adjust_readers_num(thread_num)

442 443
    @deprecated(
        since="2.0.0",
444
        update_to="paddle.distributed.InMemoryDataset._dynamic_adjust_after_train",
445
    )
446 447
    def _dynamic_adjust_after_train(self):
        if not self.is_user_set_queue_num:
448 449 450 451
            if self.use_ps_gpu:
                self.dataset.dynamic_adjust_channel_num(self.thread_num, True)
            else:
                self.dataset.dynamic_adjust_channel_num(self.thread_num, False)
452 453
        self.dataset.dynamic_adjust_readers_num(self.thread_num)

454 455 456 457
    @deprecated(
        since="2.0.0",
        update_to="paddle.distributed.InMemoryDataset._set_queue_num",
    )
J
jiaqi 已提交
458 459 460 461 462
    def set_queue_num(self, queue_num):
        """
        Set Dataset output queue num, training threads get data from queues

        Args:
463
            queue_num(int): dataset output queue num
J
jiaqi 已提交
464 465 466 467 468 469 470 471 472

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              dataset.set_queue_num(12)

        """
473
        self.is_user_set_queue_num = True
J
jiaqi 已提交
474 475
        self.queue_num = queue_num

476 477 478 479
    @deprecated(
        since="2.0.0",
        update_to="paddle.distributed.InMemoryDataset._set_parse_ins_id",
    )
480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496
    def set_parse_ins_id(self, parse_ins_id):
        """
        Set id Dataset need to parse insid

        Args:
            parse_ins_id(bool): if parse ins_id or not

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              dataset.set_parse_ins_id(True)

        """
        self.parse_ins_id = parse_ins_id

497 498
    @deprecated(
        since="2.0.0",
499 500
        update_to="paddle.distributed.InMemoryDataset._set_parse_content",
    )
501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517
    def set_parse_content(self, parse_content):
        """
        Set if Dataset need to parse content

        Args:
            parse_content(bool): if parse content or not

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              dataset.set_parse_content(True)

        """
        self.parse_content = parse_content

518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534
    def set_parse_logkey(self, parse_logkey):
        """
        Set if Dataset need to parse logkey

        Args:
            parse_content(bool): if parse logkey or not

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              dataset.set_parse_logkey(True)

        """
        self.parse_logkey = parse_logkey

535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551
    def _set_trainer_num(self, trainer_num):
        """
        Set trainer num

        Args:
            trainer_num(int): trainer num

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              dataset._set_trainer_num(1)

        """
        self.trainer_num = trainer_num

552 553 554 555
    @deprecated(
        since="2.0.0",
        update_to="paddle.distributed.InMemoryDataset._set_merge_by_sid",
    )
556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591
    def set_merge_by_sid(self, merge_by_sid):
        """
        Set if Dataset need to merge sid. If not, one ins means one Pv.

        Args:
            merge_by_sid(bool): if merge sid or not

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              dataset.set_merge_by_sid(True)

        """
        self.merge_by_sid = merge_by_sid

    def set_enable_pv_merge(self, enable_pv_merge):
        """
        Set if Dataset need to merge pv.

        Args:
            enable_pv_merge(bool): if enable_pv_merge or not

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              dataset.set_enable_pv_merge(True)

        """
        self.enable_pv_merge = enable_pv_merge

    def preprocess_instance(self):
        """
592
        Merge pv instance and convey it from input_channel to input_pv_channel.
593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644
        It will be effective when enable_pv_merge_ is True.

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.load_into_memory()
              dataset.preprocess_instance()

        """
        self.dataset.preprocess_instance()

    def set_current_phase(self, current_phase):
        """
        Set current phase in train. It is useful for untest.
        current_phase : 1 for join, 0 for update.

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.load_into_memory()
              dataset.set_current_phase(1)

        """
        self.dataset.set_current_phase(current_phase)

    def postprocess_instance(self):
        """
        Divide pv instance and convey it to input_channel.

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.load_into_memory()
              dataset.preprocess_instance()
              exe.train_from_dataset(dataset)
              dataset.postprocess_instance()

        """
        self.dataset.postprocess_instance()

645 646
    @deprecated(
        since="2.0.0",
647
        update_to="paddle.distributed.InMemoryDataset._set_fleet_send_batch_size",
648
    )
649
    def set_fleet_send_batch_size(self, fleet_send_batch_size=1024):
J
jiaqi 已提交
650
        """
651
        Set fleet send batch size, default is 1024
J
jiaqi 已提交
652 653 654 655 656 657 658 659 660 661 662 663 664

        Args:
            fleet_send_batch_size(int): fleet send batch size

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              dataset.set_fleet_send_batch_size(800)

        """
        self.fleet_send_batch_size = fleet_send_batch_size
665

666 667
    @deprecated(
        since="2.0.0",
668 669
        update_to="paddle.distributed.InMemoryDataset._set_fleet_send_sleep_seconds",
    )
670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686
    def set_fleet_send_sleep_seconds(self, fleet_send_sleep_seconds=0):
        """
        Set fleet send sleep time, default is 0

        Args:
            fleet_send_sleep_seconds(int): fleet send sleep time

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              dataset.set_fleet_send_sleep_seconds(2)

        """
        self.fleet_send_sleep_seconds = fleet_send_sleep_seconds

687 688
    @deprecated(
        since="2.0.0",
689 690
        update_to="paddle.distributed.InMemoryDataset._set_merge_by_lineid",
    )
691
    def set_merge_by_lineid(self, merge_size=2):
692 693 694 695 696
        """
        Set merge by line id, instances of same line id will be merged after
        shuffle, you should parse line id in data generator.

        Args:
697
            merge_size(int): ins size to merge. default is 2.
698 699 700 701 702 703 704 705 706

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              dataset.set_merge_by_lineid()

        """
707
        self.dataset.set_merge_by_lineid(merge_size)
708
        self.merge_by_lineid = True
709
        self.parse_ins_id = True
710

711 712
    @deprecated(
        since="2.0.0",
713 714
        update_to="paddle.distributed.InMemoryDataset._set_generate_unique_feasigns",
    )
715 716 717 718 719
    def set_generate_unique_feasigns(self, generate_uni_feasigns, shard_num):
        self.dataset.set_generate_unique_feasigns(generate_uni_feasigns)
        self.gen_uni_feasigns = generate_uni_feasigns
        self.local_shard_num = shard_num

720 721
    @deprecated(
        since="2.0.0",
722 723 724 725 726 727 728 729
        update_to="paddle.distributed.InMemoryDataset._generate_local_tables_unlock",
    )
    def generate_local_tables_unlock(
        self, table_id, fea_dim, read_thread_num, consume_thread_num, shard_num
    ):
        self.dataset.generate_local_tables_unlock(
            table_id, fea_dim, read_thread_num, consume_thread_num, shard_num
        )
730

731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753
    def set_date(self, date):
        """
        :api_attr: Static Graph

        Set training date for pull sparse parameters, saving and loading model. Only used in psgpu

        Args:
            date(str): training date(format : YYMMDD). eg.20211111

        Examples:
            .. code-block:: python

                import paddle.fluid as fluid

                dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
                dataset.set_date("20211111")
        """
        year = int(date[:4])
        month = int(date[4:6])
        day = int(date[6:])
        if self.use_ps_gpu and core._is_compiled_with_heterps():
            self.psgpu.set_date(year, month, day)

754 755 756 757
    @deprecated(
        since="2.0.0",
        update_to="paddle.distributed.InMemoryDataset.load_into_memory",
    )
758
    def load_into_memory(self, is_shuffle=False):
759 760 761
        """
        Load data into memory

762 763 764
         Args:
            is_shuffle(bool): whether to use local shuffle, default is False

765 766 767
        Examples:
            .. code-block:: python

768
              # required: skiptest
769 770 771 772 773
              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.load_into_memory()
774
        """
775
        self._prepare_to_run()
776 777 778 779 780
        if not self.use_ps_gpu:
            self.dataset.load_into_memory()
        elif core._is_compiled_with_heterps():
            self.psgpu.set_dataset(self.dataset)
            self.psgpu.load_into_memory(is_shuffle)
D
dongdaxiang 已提交
781

782 783
    @deprecated(
        since="2.0.0",
784 785
        update_to="paddle.distributed.InMemoryDataset.preload_into_memory",
    )
786
    def preload_into_memory(self, thread_num=None):
J
jiaqi 已提交
787 788 789
        """
        Load data into memory in async mode

790 791 792
        Args:
            thread_num(int): preload thread num

J
jiaqi 已提交
793 794 795
        Examples:
            .. code-block:: python

796
              # required: skiptest
J
jiaqi 已提交
797 798 799 800 801 802 803 804
              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.preload_into_memory()
              dataset.wait_preload_done()
        """
        self._prepare_to_run()
805 806 807 808
        if thread_num is None:
            thread_num = self.thread_num
        self.dataset.set_preload_thread_num(thread_num)
        self.dataset.create_preload_readers()
J
jiaqi 已提交
809 810
        self.dataset.preload_into_memory()

811 812 813 814
    @deprecated(
        since="2.0.0",
        update_to="paddle.distributed.InMemoryDataset.wait_preload_done",
    )
J
jiaqi 已提交
815 816 817 818 819 820 821
    def wait_preload_done(self):
        """
        Wait preload_into_memory done

        Examples:
            .. code-block:: python

822
              # required: skiptest
J
jiaqi 已提交
823 824 825 826 827 828 829 830
              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.preload_into_memory()
              dataset.wait_preload_done()
        """
        self.dataset.wait_preload_done()
831
        self.dataset.destroy_preload_readers()
J
jiaqi 已提交
832

833 834 835 836
    @deprecated(
        since="2.0.0",
        update_to="paddle.distributed.InMemoryDataset.local_shuffle",
    )
D
dongdaxiang 已提交
837
    def local_shuffle(self):
838 839 840
        """
        Local shuffle

841 842 843
        Examples:
            .. code-block:: python

844
              # required: skiptest
845 846 847 848 849 850
              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.load_into_memory()
              dataset.local_shuffle()
851
        """
852
        self.dataset.local_shuffle()
D
dongdaxiang 已提交
853

854 855 856 857
    @deprecated(
        since="2.0.0",
        update_to="paddle.distributed.InMemoryDataset.global_shuffle",
    )
858
    def global_shuffle(self, fleet=None, thread_num=12):
859 860
        """
        Global shuffle.
861 862 863
        Global shuffle can be used only in distributed mode. i.e. multiple
        processes on single machine or multiple machines training together.
        If you run in distributed mode, you should pass fleet instead of None.
864

865
        Examples:
866 867
            .. code-block:: python

868
              # required: skiptest
869 870 871 872 873 874 875
              import paddle.fluid as fluid
              from paddle.fluid.incubate.fleet.parameter_server.pslib import fleet
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.load_into_memory()
              dataset.global_shuffle(fleet)
876 877

        Args:
878
            fleet(Fleet): fleet singleton. Default None.
879
            thread_num(int): shuffle thread num. Default is 12.
880

881
        """
882
        if fleet is not None:
883 884
            if hasattr(fleet, "barrier_worker"):
                print("pscore fleet")
885 886 887
                fleet.barrier_worker()
            else:
                fleet._role_maker.barrier_worker()
888 889
            if self.trainer_num == -1:
                self.trainer_num = fleet.worker_num()
890
        if self.fleet_send_batch_size is None:
891 892 893
            self.fleet_send_batch_size = 1024
        if self.fleet_send_sleep_seconds is None:
            self.fleet_send_sleep_seconds = 0
894
        self.dataset.register_client2client_msg_handler()
895
        self.dataset.set_trainer_num(self.trainer_num)
J
jiaqi 已提交
896
        self.dataset.set_fleet_send_batch_size(self.fleet_send_batch_size)
897
        self.dataset.set_fleet_send_sleep_seconds(self.fleet_send_sleep_seconds)
898
        if fleet is not None:
899
            if hasattr(fleet, "barrier_worker"):
900 901 902
                fleet.barrier_worker()
            else:
                fleet._role_maker.barrier_worker()
903
        self.dataset.global_shuffle(thread_num)
904
        if fleet is not None:
905
            if hasattr(fleet, "barrier_worker"):
906 907 908
                fleet.barrier_worker()
            else:
                fleet._role_maker.barrier_worker()
909 910 911
        if self.merge_by_lineid:
            self.dataset.merge_by_lineid()
        if fleet is not None:
912
            if hasattr(fleet, "barrier_worker"):
913 914 915
                fleet.barrier_worker()
            else:
                fleet._role_maker.barrier_worker()
D
dongdaxiang 已提交
916

917 918 919 920
    @deprecated(
        since="2.0.0",
        update_to="paddle.distributed.InMemoryDataset.release_memory",
    )
921 922
    def release_memory(self):
        """
923
        :api_attr: Static Graph
924

925 926
        Release InMemoryDataset memory data, when data will not be used again.

927 928 929
        Examples:
            .. code-block:: python

930
              # required: skiptest
931 932 933 934 935 936 937 938 939 940 941 942
              import paddle.fluid as fluid
              from paddle.fluid.incubate.fleet.parameter_server.pslib import fleet
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.load_into_memory()
              dataset.global_shuffle(fleet)
              exe = fluid.Executor(fluid.CPUPlace())
              exe.run(fluid.default_startup_program())
              exe.train_from_dataset(fluid.default_main_program(), dataset)
              dataset.release_memory()

943 944
        """
        self.dataset.release_memory()
D
dongdaxiang 已提交
945

946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969
    def get_pv_data_size(self):
        """
        Get memory data size of Pv, user can call this function to know the pv num
        of ins in all workers after load into memory.

        Note:
            This function may cause bad performance, because it has barrier

        Returns:
            The size of memory pv data.

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.load_into_memory()
              print dataset.get_pv_data_size()

        """
        return self.dataset.get_pv_data_size()

970 971
    @deprecated(
        since="2.0.0",
972 973
        update_to="paddle.distributed.InMemoryDataset.get_memory_data_size",
    )
974 975 976 977 978 979 980 981 982 983 984 985 986 987
    def get_memory_data_size(self, fleet=None):
        """
        Get memory data size, user can call this function to know the num
        of ins in all workers after load into memory.

        Note:
            This function may cause bad performance, because it has barrier

        Args:
            fleet(Fleet): Fleet Object.

        Returns:
            The size of memory data.

988 989 990
        Examples:
            .. code-block:: python

991
              # required: skiptest
992 993 994 995 996 997 998
              import paddle.fluid as fluid
              from paddle.fluid.incubate.fleet.parameter_server.pslib import fleet
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.load_into_memory()
              print dataset.get_memory_data_size(fleet)
999 1000 1001

        """
        import numpy as np
1002

1003 1004 1005 1006
        local_data_size = self.dataset.get_memory_data_size()
        local_data_size = np.array([local_data_size])
        if fleet is not None:
            global_data_size = local_data_size * 0
1007 1008 1009
            fleet._role_maker.all_reduce_worker(
                local_data_size, global_data_size
            )
1010 1011 1012
            return global_data_size[0]
        return local_data_size[0]

1013 1014
    @deprecated(
        since="2.0.0",
1015 1016
        update_to="paddle.distributed.InMemoryDataset.get_shuffle_data_size",
    )
1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031
    def get_shuffle_data_size(self, fleet=None):
        """
        Get shuffle data size, user can call this function to know the num
        of ins in all workers after local/global shuffle.

        Note:
            This function may cause bad performance to local shuffle,
            because it has barrier. It does not affect global shuffle.

        Args:
            fleet(Fleet): Fleet Object.

        Returns:
            The size of shuffle data.

1032 1033 1034
        Examples:
            .. code-block:: python

1035
              # required: skiptest
1036 1037 1038 1039 1040 1041 1042 1043
              import paddle.fluid as fluid
              from paddle.fluid.incubate.fleet.parameter_server.pslib import fleet
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.load_into_memory()
              dataset.global_shuffle(fleet)
              print dataset.get_shuffle_data_size(fleet)
1044 1045 1046

        """
        import numpy as np
1047

1048 1049
        local_data_size = self.dataset.get_shuffle_data_size()
        local_data_size = np.array([local_data_size])
1050
        print('global shuffle local_data_size: ', local_data_size)
1051 1052
        if fleet is not None:
            global_data_size = local_data_size * 0
1053
            if hasattr(fleet, "util"):
1054 1055
                global_data_size = fleet.util.all_reduce(local_data_size)
            else:
1056 1057 1058
                fleet._role_maker.all_reduce_worker(
                    local_data_size, global_data_size
                )
1059 1060 1061
            return global_data_size[0]
        return local_data_size[0]

Y
yaoxuefeng 已提交
1062 1063 1064 1065 1066 1067 1068
    def _set_heter_ps(self, enable_heter_ps=False):
        """
        Set heter ps mode
        user no need to call this function.
        """
        self.dataset.set_heter_ps(enable_heter_ps)

D
danleifeng 已提交
1069 1070
    def set_graph_config(self, config):
        """
1071
        Set graph config, user can set graph config in gpu graph mode.
D
danleifeng 已提交
1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101

        Args:
            config(dict): config dict.

        Returns:
            The size of shuffle data.

        Examples:
            .. code-block:: python

              # required: skiptest
              import paddle.fluid as fluid
              from paddle.fluid.incubate.fleet.parameter_server.pslib import fleet
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              graph_config = {"walk_len": 24,
                    "walk_degree": 10,
                    "once_sample_startid_len": 80000,
                    "sample_times_one_chunk": 5,
                    "window": 3,
                    "debug_mode": 0,
                    "batch_size": 800,
                    "meta_path": "cuid2clk-clk2cuid;cuid2conv-conv2cuid;clk2cuid-cuid2clk;clk2cuid-cuid2conv",
                    "gpu_graph_training": 1}
              dataset.set_graph_config(graph_config)

        """
        self.proto_desc.graph_config.walk_degree = config.get("walk_degree", 1)
        self.proto_desc.graph_config.walk_len = config.get("walk_len", 20)
        self.proto_desc.graph_config.window = config.get("window", 5)
        self.proto_desc.graph_config.once_sample_startid_len = config.get(
1102 1103
            "once_sample_startid_len", 8000
        )
D
danleifeng 已提交
1104
        self.proto_desc.graph_config.sample_times_one_chunk = config.get(
1105 1106
            "sample_times_one_chunk", 10
        )
D
danleifeng 已提交
1107 1108 1109
        self.proto_desc.graph_config.batch_size = config.get("batch_size", 1)
        self.proto_desc.graph_config.debug_mode = config.get("debug_mode", 0)
        self.proto_desc.graph_config.first_node_type = config.get(
1110 1111
            "first_node_type", ""
        )
D
danleifeng 已提交
1112 1113
        self.proto_desc.graph_config.meta_path = config.get("meta_path", "")
        self.proto_desc.graph_config.gpu_graph_training = config.get(
1114 1115
            "gpu_graph_training", True
        )
L
lxsbupt 已提交
1116 1117 1118 1119 1120 1121 1122 1123
        self.proto_desc.graph_config.sage_mode = config.get("sage_mode", False)
        self.proto_desc.graph_config.samples = config.get("samples", "")
        self.proto_desc.graph_config.train_table_cap = config.get(
            "train_table_cap", 800000
        )
        self.proto_desc.graph_config.infer_table_cap = config.get(
            "infer_table_cap", 800000
        )
D
danleifeng 已提交
1124 1125
        self.dataset.set_gpu_graph_mode(True)

L
lxsbupt 已提交
1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151 1152 1153 1154 1155 1156 1157 1158 1159
    def set_pass_id(self, pass_id):
        """
        Set pass id, user can set pass id in gpu graph mode.

        Args:
            pass_id: pass id.

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              pass_id = 0
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              dataset.set_pass_id(pass_id)
        """
        self.pass_id = pass_id
        self.dataset.set_pass_id(pass_id)

    def get_pass_id(self):
        """
        Get pass id, user can set pass id in gpu graph mode.

        Returns:
            The pass id.

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
              pass_id = dataset.get_pass_id()
        """
        return self.pass_id

X
xjqbest 已提交
1160

D
dongdaxiang 已提交
1161
class QueueDataset(DatasetBase):
1162 1163 1164
    """
    QueueDataset, it will process data streamly.

1165 1166 1167 1168 1169 1170
    Examples:
        .. code-block:: python

          import paddle.fluid as fluid
          dataset = fluid.DatasetFactory().create_dataset("QueueDataset")

1171
    """
D
dongdaxiang 已提交
1172

D
dongdaxiang 已提交
1173
    def __init__(self):
1174
        """
D
dongdaxiang 已提交
1175 1176
        Initialize QueueDataset
        This class should be created by DatasetFactory
1177
        """
1178
        super().__init__()
D
dongdaxiang 已提交
1179
        self.proto_desc.name = "MultiSlotDataFeed"
X
xujiaqi01 已提交
1180

1181 1182 1183 1184
    @deprecated(
        since="2.0.0",
        update_to="paddle.distributed.QueueDataset._prepare_to_run",
    )
1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198
    def _prepare_to_run(self):
        """
        Set data_feed_desc/thread num/filelist before run,
        user no need to call this function.
        """
        if self.thread_num > len(self.filelist):
            self.thread_num = len(self.filelist)
        if self.thread_num == 0:
            self.thread_num = 1
        self.dataset.set_thread_num(self.thread_num)
        self.dataset.set_filelist(self.filelist)
        self.dataset.set_data_feed_desc(self.desc())
        self.dataset.create_readers()

X
xujiaqi01 已提交
1199
    def local_shuffle(self):
1200
        """
1201
        Local shuffle data.
D
dongdaxiang 已提交
1202

D
dongdaxiang 已提交
1203 1204
        Local shuffle is not supported in QueueDataset
        NotImplementedError will be raised
1205 1206 1207 1208 1209 1210 1211 1212

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("QueueDataset")
              dataset.local_shuffle()

1213 1214 1215
        Raises:
            NotImplementedError: QueueDataset does not support local shuffle

1216
        """
D
dongdaxiang 已提交
1217 1218
        raise NotImplementedError(
            "QueueDataset does not support local shuffle, "
1219 1220
            "please use InMemoryDataset for local_shuffle"
        )
X
xujiaqi01 已提交
1221

1222
    def global_shuffle(self, fleet=None):
1223
        """
1224 1225
        Global shuffle data.

D
dongdaxiang 已提交
1226 1227
        Global shuffle is not supported in QueueDataset
        NotImplementedError will be raised
1228

1229 1230 1231
        Args:
            fleet(Fleet): fleet singleton. Default None.

1232 1233 1234 1235 1236 1237 1238 1239
        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              from paddle.fluid.incubate.fleet.parameter_server.pslib import fleet
              dataset = fluid.DatasetFactory().create_dataset("QueueDataset")
              dataset.global_shuffle(fleet)

1240 1241 1242
        Raises:
            NotImplementedError: QueueDataset does not support global shuffle

1243
        """
D
dongdaxiang 已提交
1244 1245
        raise NotImplementedError(
            "QueueDataset does not support global shuffle, "
1246 1247
            "please use InMemoryDataset for global_shuffle"
        )
H
hutuxian 已提交
1248 1249 1250 1251 1252


class FileInstantDataset(DatasetBase):
    """
    FileInstantDataset, it will process data streamly.
1253 1254 1255 1256 1257 1258

    Examples:
        .. code-block:: python

          import paddle.fluid as fluid
          dataset = fluid.DatasetFactory.create_dataset("FileInstantDataset")
H
hutuxian 已提交
1259 1260 1261 1262
    """

    def __init__(self):
        """
1263 1264
        Initialize FileInstantDataset
        This class should be created by DatasetFactory
H
hutuxian 已提交
1265
        """
1266
        super().__init__()
H
hutuxian 已提交
1267 1268 1269 1270
        self.proto_desc.name = "MultiSlotFileInstantDataFeed"

    def local_shuffle(self):
        """
1271 1272
        Local shuffle
        FileInstantDataset does not support local shuffle
H
hutuxian 已提交
1273 1274 1275
        """
        raise NotImplementedError(
            "FileInstantDataset does not support local shuffle, "
1276 1277
            "please use InMemoryDataset for local_shuffle"
        )
H
hutuxian 已提交
1278 1279 1280 1281

    def global_shuffle(self, fleet=None):
        """
        Global shuffle
1282
        FileInstantDataset does not support global shuffle
H
hutuxian 已提交
1283 1284 1285
        """
        raise NotImplementedError(
            "FileInstantDataset does not support global shuffle, "
1286 1287
            "please use InMemoryDataset for global_shuffle"
        )
H
hutuxian 已提交
1288 1289 1290 1291 1292 1293 1294 1295 1296 1297


class BoxPSDataset(InMemoryDataset):
    """
    BoxPSDataset: derived from InMemoryDataset.

    Examples:
        .. code-block:: python

          import paddle.fluid as fluid
H
hutuxian 已提交
1298
          dataset = fluid.DatasetFactory().create_dataset("BoxPSDataset")
H
hutuxian 已提交
1299 1300 1301 1302
    """

    def __init__(self):
        """
1303 1304
        Initialize BoxPSDataset
        This class should be created by DatasetFactory
H
hutuxian 已提交
1305
        """
1306
        super().__init__()
H
hutuxian 已提交
1307
        self.boxps = core.BoxPS(self.dataset)
1308
        self.proto_desc.name = "PaddleBoxDataFeed"
H
hutuxian 已提交
1309

H
hutuxian 已提交
1310 1311 1312 1313 1314 1315 1316 1317 1318
    def set_date(self, date):
        """
        Workaround for date
        """
        year = int(date[:4])
        month = int(date[4:6])
        day = int(date[6:])
        self.boxps.set_date(year, month, day)

H
hutuxian 已提交
1319 1320
    def begin_pass(self):
        """
1321
        Begin Pass
1322
        Notify BoxPS to load sparse parameters of next pass to GPU Memory
H
hutuxian 已提交
1323 1324 1325 1326 1327 1328 1329 1330

        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("BoxPSDataset")
              dataset.begin_pass()
        """
H
hutuxian 已提交
1331 1332
        self.boxps.begin_pass()

1333
    def end_pass(self, need_save_delta):
H
hutuxian 已提交
1334
        """
1335
        End Pass
1336
        Notify BoxPS that current pass ended
H
hutuxian 已提交
1337 1338 1339 1340 1341
        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("BoxPSDataset")
1342
              dataset.end_pass(True)
H
hutuxian 已提交
1343
        """
1344
        self.boxps.end_pass(need_save_delta)
H
hutuxian 已提交
1345 1346 1347

    def wait_preload_done(self):
        """
T
tianshuo78520a 已提交
1348
        Wait async preload done
1349
        Wait Until Feed Pass Done
H
hutuxian 已提交
1350 1351 1352 1353 1354 1355 1356 1357 1358 1359
        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("BoxPSDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.preload_into_memory()
              dataset.wait_preload_done()
        """
H
hutuxian 已提交
1360 1361 1362 1363
        self.boxps.wait_feed_pass_done()

    def load_into_memory(self):
        """
H
hutuxian 已提交
1364 1365 1366 1367 1368 1369 1370 1371 1372
        Load next pass into memory and notify boxps to fetch its emb from SSD
        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("BoxPSDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.load_into_memory()
1373
        """
H
hutuxian 已提交
1374 1375 1376 1377 1378
        self._prepare_to_run()
        self.boxps.load_into_memory()

    def preload_into_memory(self):
        """
H
hutuxian 已提交
1379 1380 1381 1382 1383 1384 1385 1386 1387 1388
        Begin async preload next pass while current pass may be training
        Examples:
            .. code-block:: python

              import paddle.fluid as fluid
              dataset = fluid.DatasetFactory().create_dataset("BoxPSDataset")
              filelist = ["a.txt", "b.txt"]
              dataset.set_filelist(filelist)
              dataset.preload_into_memory()
        """
H
hutuxian 已提交
1389 1390
        self._prepare_to_run()
        self.boxps.preload_into_memory()
H
hutuxian 已提交
1391 1392 1393 1394 1395

    def _dynamic_adjust_before_train(self, thread_num):
        if not self.is_user_set_queue_num:
            self.dataset.dynamic_adjust_channel_num(thread_num, True)
        self.dataset.dynamic_adjust_readers_num(thread_num)
1396 1397 1398

    def _dynamic_adjust_after_train(self):
        pass
1399 1400 1401

    def slots_shuffle(self, slots):
        """
1402 1403
        Slots Shuffle
        Slots Shuffle is a shuffle method in slots level, which is usually used
1404
        in sparse feature with large scale of instances. To compare the metric, i.e.
1405
        auc while doing slots shuffle on one or several slots with baseline to
1406
        evaluate the importance level of slots(features).
1407

1408 1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 1419
        Args:
            slots(list[string]): the set of slots(string) to do slots shuffle.

        Examples:
            import paddle.fluid as fluid
            dataset = fluid.DatasetFactory().create_dataset("InMemoryDataset")
            dataset.set_merge_by_lineid()
            #suppose there is a slot 0
            dataset.slots_shuffle(['0'])
        """
        slots_set = set(slots)
        self.boxps.slots_shuffle(slots_set)