wehub-resource-sync
89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
# Copyright (c) 2026 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.
|
|
from __future__ import annotations
|
|
|
|
import warnings
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from paddle.io import DataLoader as PaddleDataLoader
|
|
|
|
from ._utils.collate import (
|
|
default_collate as default_collate,
|
|
)
|
|
from ._utils.worker import (
|
|
get_worker_info as get_worker_info,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
from paddle.io.dataloader import BatchSampler
|
|
from paddle.io.dataloader.dataset import Dataset
|
|
from paddle.io.reader import _CollateFn
|
|
|
|
|
|
class DataLoader(PaddleDataLoader):
|
|
def __init__(
|
|
self,
|
|
dataset: Dataset[Any],
|
|
batch_size: int | None = 1,
|
|
shuffle: bool = False,
|
|
sampler: BatchSampler | None = None,
|
|
batch_sampler: BatchSampler | None = None,
|
|
num_workers: int = 0,
|
|
collate_fn: _CollateFn | None = None,
|
|
pin_memory: bool = False,
|
|
drop_last: bool = False,
|
|
timeout: float = 0,
|
|
worker_init_fn: Callable[[int], None] | None = None,
|
|
multiprocessing_context=None,
|
|
generator=None,
|
|
*,
|
|
prefetch_factor: int | None = None,
|
|
persistent_workers: bool = False,
|
|
pin_memory_device: str = "",
|
|
in_order: bool = True,
|
|
) -> None:
|
|
if (
|
|
pin_memory is True
|
|
or multiprocessing_context is not None
|
|
or generator is not None
|
|
or prefetch_factor is not None
|
|
or len(pin_memory_device) > 0
|
|
or in_order is False
|
|
):
|
|
warnings.warn(
|
|
"pin_memory, multiprocessing_context, generator, prefetch_factor, pin_memory_device, in_order are currently not supported in DataLoader and will be ignored."
|
|
)
|
|
|
|
if sampler is not None:
|
|
if batch_sampler is not None:
|
|
raise ValueError(
|
|
"Cannot specify both sampler and batch_sampler"
|
|
)
|
|
batch_sampler = sampler
|
|
|
|
super().__init__(
|
|
dataset=dataset,
|
|
batch_sampler=batch_sampler,
|
|
batch_size=batch_size,
|
|
shuffle=shuffle,
|
|
drop_last=drop_last,
|
|
collate_fn=collate_fn,
|
|
num_workers=num_workers,
|
|
timeout=timeout,
|
|
worker_init_fn=worker_init_fn,
|
|
persistent_workers=persistent_workers,
|
|
)
|