unienv_data.samplers.multiprocessing_sampler¶
MultiProcessingSampleManager(target_fn: Callable[[TaskInfoT], ResultT], n_workers: int = 4, ctx: Optional[BaseContext] = None, done_event=None, doze_time: float = 0.001, daemon: Optional[bool] = None, worker_init_fn: Optional[Callable[[int], None]] = None, worker_seeds: Optional[Sequence[int]] = None)
MultiprocessingSampler(sampler: BatchSampler[SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType, BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], n_workers: int = 4, n_buffers: int = 8, ctx: Optional[BaseContext] = None, doze_time: float = 0.005, daemon: Optional[bool] = None, shared_memory: bool = False)
Bases: BatchSampler[SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType, BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Run sampler's sample* functions (and epoch iteration over its get*_at
functions) in n_workers background processes, keeping n_buffers batches queued.
Workers pay off when producing a batch is expensive (decoding, transformations,
storage reads). For data already in RAM (e.g. an in-memory ReplayBuffer),
sampling in-process is usually faster (e.g. 0.7 ms vs 3 ms for 256 samples of 48 KiB
with 4 workers): every batch is otherwise copied to the parent through shared memory.
Workers run their torch ops single-threaded (see _limit_worker_intraop_threads).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
shared_memory
|
bool
|
For |
False
|
set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None
Overwrite existing samples using flattened data.
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
extend_flattened(value: BArrayType) -> None
Append a batched block of flattened samples.
set_at(idx: Union[IndexableType, BArrayType], value: BatchT) -> None
Overwrite existing samples using structured data.
remove_at(idx: Union[IndexableType, BArrayType]) -> None
Remove one or more samples from the batch.
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.
get_column(nested_keys: Sequence[str]) -> BatchBase[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a nested field inside each sample.
sample_flat_with_metadata() -> Tuple[SamplerArrayType, Optional[Dict[str, Any]]]
Sample a flattened batch together with optional metadata.
sample_with_metadata() -> Tuple[SamplerBatchT, Optional[Dict[str, Any]]]
Sample a structured batch together with optional metadata.
epoch_iter() -> Iterator[SamplerBatchT]
Iterate once over a random permutation of the source data.
epoch_iter_with_metadata() -> Iterator[Tuple[SamplerBatchT, Optional[Dict[str, Any]]]]
Iterate once over shuffled structured batches plus metadata.
epoch_flat_iter() -> Iterator[SamplerArrayType]
Iterate once over shuffled batches in flattened form.
epoch_flat_iter_with_metadata() -> Iterator[Tuple[SamplerArrayType, Optional[Dict[str, Any]]]]
Iterate once over shuffled flattened batches plus metadata.
worker_loop(fetch_fn: Callable[[BatchSampler[SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType, BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], Any], Any], workid_queue: Queue, result_queue: Queue, done_event: Event, doze_time: float = 0.005, init_fn: Optional[Callable[[int], None]] = None, seed: Optional[int] = None)
wrap_sample_function(sampler: BatchSampler[SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType, BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], n_workers: int, n_buffers: int, fn: Callable[[], Any], ctx: Optional[BaseContext] = None, doze_time: float = 0.001, daemon: Optional[bool] = None, done_event=None, exemplar_fn: Optional[Callable[[], Any]] = None)
exemplar_fn (optional) returns a result-shaped exemplar; its torch CPU tensors get
pre-allocated shared-memory slots that workers write their results into (see
:class:_SharedSlotTask).
wrap_epoch_iter_function(sampler: BatchSampler[SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType, BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], n_workers: int, n_buffers: int, getter_fn: Callable[[Any], Any], ctx: Optional[BaseContext] = None, doze_time: float = 0.005, daemon: Optional[bool] = None, done_event=None)