网站模板设计教程,建站系统主要包括什么,wordpress 万能搜索页,伪春菜wordpressDataLoader 类
功能与作用 PyTorch 是一个流行的开源机器学习库#xff0c;它提供了一个名为 DataLoader 的类#xff0c;用于加载数据集并将其封装成一个可迭代的对象。DataLoader 可以自动地将数据集划分为多个批次#xff0c;并在训练过程中迭代地返回这些批次。是用于加…DataLoader 类
功能与作用 PyTorch 是一个流行的开源机器学习库它提供了一个名为 DataLoader 的类用于加载数据集并将其封装成一个可迭代的对象。DataLoader 可以自动地将数据集划分为多个批次并在训练过程中迭代地返回这些批次。是用于加载数据集的重要组件。DataLoader 结合了数据集和采样器提供了一个可迭代的数据加载器支持多进程数据加载、自定义加载顺序和可选的自动批处理collation以及内存锁定。 DataLoader 类的一些关键特性 批量加载Batching可以指定每个批次的大小。多进程加载Multiprocessing可以利用多个进程来加速数据的加载。数据打乱Shuffling可以在每个epoch开始时随机打乱数据以提高模型训练的泛化能力。数据采样Sampling可以自定义数据采样方式例如可以定义一个采样器来选择数据集中的特定样本。数据预处理Data Preprocessing可以在加载数据时应用预处理操作例如数据增强。
基本结构介绍 参数介绍 dataset (Dataset): 要加载数据的数据集。batch_size (int, optional): 每个批次加载的样本数默认为 1。shuffle (bool, optional): 是否在每个 epoch 打乱数据默认为 False。sampler (Sampler or Iterable, optional): 定义从数据集中抽取样本的策略。batch_sampler (Sampler or Iterable, optional): 类似于 sampler但每次返回一批索引。num_workers (int, optional): 用于数据加载的子进程数默认为 0即在主进程中加载数据。collate_fn (Callable, optional): 合并样本以形成 mini-batch 的函数。pin_memory (bool, optional): 是否将数据复制到 CUDA 固定内存中默认为 False。drop_last (bool, optional): 是否丢弃最后一个不完整的批次默认为 False。timeout (numeric, optional): 从工作进程收集一批数据的超时时间默认为 0。worker_init_fn (Callable, optional): 每个工作进程启动时调用的函数。multiprocessing_context (str or multiprocessing.context.BaseContext, optional): 多进程上下文。generator (torch.Generator, optional): 用于生成随机索引的随机数生成器。prefetch_factor (int, optional): 每个工作进程提前加载的批次数。persistent_workers (bool, optional): 是否在数据集消费完毕后保持工作进程活动默认为 False。pin_memory_device (str, optional): 将数据固定到的设备。 方法介绍 __init__: 初始化 DataLoader 实例。_get_iterator: 返回一个迭代器用于从 DataLoader 中获取数据。multiprocessing_context: 获取或设置多进程上下文。__setattr__: 设置实例属性如果 DataLoader 已经初始化则不能修改某些属性。__iter__: 返回一个迭代器用于遍历 DataLoader。_auto_collation: 属性指示是否启用自动批处理。_index_sampler: 属性获取用于生成索引的实际采样器。__len__: 返回 DataLoader 的长度即可以生成的批次总数。check_worker_number_rationality: 检查工作进程数是否合理基于系统资源。 其他相关类 _BaseDataLoaderIter 类这是 DataLoader 的基类迭代器提供了多进程数据加载的底层逻辑。_SingleProcessDataLoaderIter 类这是 DataLoader 的单进程迭代器用于在主进程中加载数据不涉及多进程通信。_MultiProcessingDataLoaderIter 类这是 DataLoader 的多进程迭代器用于在多个进程中加载数据提高了数据加载的效率。辅助函数和类例如 default_collate、default_convert、get_worker_info 等这些函数和类提供了数据加载过程中需要的一些辅助功能。
DataLoader 和 Dataset
DataLoader 和 Dataset 是 PyTorch 中用于数据加载和预处理的两个非常重要的类它们在数据加载流程中扮演着不同的角色但又紧密相关。 Dataset Dataset 是 PyTorch 中的一个抽象类它定义了数据集的接口规范。具体来说Dataset 类需要实现以下两个方法 __len__: 返回数据集中样本的总数。__getitem__: 接受一个索引值并返回对应的样本。 Dataset 的子类可以是任何类型的数据集比如图像、文本或音频数据集。开发者需要根据具体的数据类型和格式继承 Dataset 类并实现这两个方法。这样Dataset 的实例就可以被看作是一个包含所有数据的容器。 DataLoader DataLoader 是 PyTorch 中用于加载数据集的工具。它提供了一种简便的方式来迭代地访问 Dataset 中的数据。DataLoader 支持多进程数据加载可以显著提高数据读取和预处理的效率。DataLoader 的主要特点包括 多进程数据加载通过 num_workers 参数可以指定多个子进程来并行加载数据这样可以有效地利用多核 CPU 资源加速数据的准备过程。批处理DataLoader 可以将多个样本组合成一个批次batch并自动处理数据的打乱、采样等操作。可定制的数据预处理通过 collate_fn 参数用户可以定义自己的数据预处理函数以适应特定的需求。 区别与联系 区别 Dataset 是一个抽象类定义了数据集的结构和访问方式而 DataLoader 是一个具体类用于从 Dataset 中加载和迭代数据。Dataset 负责存储数据和提供单个样本的访问DataLoader 负责按批次加载数据并支持多进程加载和批处理。Dataset 的实现侧重于数据的组织和读取而 DataLoader 的实现侧重于数据加载的效率和灵活性。 联系 DataLoader 依赖于 Dataset它需要一个 Dataset 的实例作为数据源。DataLoader 通过 Dataset 提供的样本来创建批次数据它封装了 Dataset 的访问方式使得数据的迭代更加高效和方便。在使用时通常先定义一个 Dataset 的子类来管理特定的数据集然后创建 DataLoader 的实例来加载和迭代这个数据集。
总的来说Dataset 和 DataLoader 是 PyTorch 数据处理流程中不可或缺的两个组件它们共同工作为模型训练提供了高效、灵活的数据支持。
示例说明
import torch
from torch.utils.data import Dataset, DataLoader# 定义一个简单的数据集
class SimpleDataset(Dataset):def __init__(self, data, labels):self.data dataself.labels labelsdef __len__(self):return len(self.data)def __getitem__(self, index):return self.data[index], self.labels[index]# 创建数据和标签
data torch.randn(100, 10) # 100个样本每个样本10个特征
labels torch.randint(0, 2, (100,)) # 100个样本的二分类标签# 实例化数据集
dataset SimpleDataset(data, labels)# 实例化DataLoader
data_loader DataLoader(dataset, batch_size4, shuffleTrue, num_workers2)# 遍历DataLoader
for batch_idx, (inputs, targets) in enumerate(data_loader):print(fBatch {batch_idx1}:)print(fInputs: {inputs})print(fTargets: {targets})# 这里可以加入模型训练的代码相关源码
PyTorch源码地址https://github.com/pytorch/pytorch/blob/main/torch/utils/data/dataloader.py
# mypy: allow-untyped-defs
rDefinition of the DataLoader and associated iterators that subclass _BaseDataLoaderIter.To support these two classes, in ./_utils we define many utility methods and
functions to be run in multiprocessing. E.g., the data loading worker loop is
in ./_utils/worker.py.
import functools
import itertools
import logging
import multiprocessing as python_multiprocessing
import os
import queue
import threading
import warnings
from typing import Any, Callable, Generic, Iterable, List, Optional, TypeVar, Unionimport torch
import torch.distributed as dist
import torch.utils.data.graph_settings
from torch._utils import ExceptionWrapper
from torch.utils.data import _utils
from torch.utils.data.datapipes.datapipe import (_IterDataPipeSerializationWrapper,_MapDataPipeSerializationWrapper,IterDataPipe,MapDataPipe,
)
from torch.utils.data.dataset import Dataset, IterableDataset
from torch.utils.data.sampler import (BatchSampler,RandomSampler,Sampler,SequentialSampler,
)__all__ [DataLoader,get_worker_info,default_collate,default_convert,
]_T TypeVar(_T)
_T_co TypeVar(_T_co, covariantTrue)
_worker_init_fn_t Callable[[int], None]# Ideally we would parameterize DataLoader by the return type of collate_fn, but there is currently no way to have that
# type parameter set to a default value if the user doesnt pass in a custom collate_fn.
# See https://github.com/python/mypy/issues/3737.
_collate_fn_t Callable[[List[_T]], Any]# These functions used to be defined in this file. However, it was moved to
# _utils/collate.py. Although it is rather hard to access this from user land
# (one has to explicitly directly import torch.utils.data.dataloader), there
# probably is user code out there using it. This aliasing maintains BC in this
# aspect.
default_collate: _collate_fn_t _utils.collate.default_collate
default_convert _utils.collate.default_convertget_worker_info _utils.worker.get_worker_infologger logging.getLogger(__name__)class _DatasetKind:Map 0Iterable 1staticmethoddef create_fetcher(kind, dataset, auto_collation, collate_fn, drop_last):if kind _DatasetKind.Map:return _utils.fetch._MapDatasetFetcher(dataset, auto_collation, collate_fn, drop_last)else:return _utils.fetch._IterableDatasetFetcher(dataset, auto_collation, collate_fn, drop_last)class _InfiniteConstantSampler(Sampler):rAnalogous to itertools.repeat(None, None).Used as sampler for :class:~torch.utils.data.IterableDataset.def __iter__(self):while True:yield Nonedef _get_distributed_settings():if dist.is_available() and dist.is_initialized():return dist.get_world_size(), dist.get_rank()else:return 1, 0def _sharding_worker_init_fn(worker_init_fn, world_size, rank_id, worker_id):global_worker_id worker_idinfo torch.utils.data.get_worker_info()assert info is not Nonetotal_workers info.num_workersdatapipe info.datasetassert isinstance(datapipe, (IterDataPipe, MapDataPipe))# To distribute elements across distributed process evenly, we should shard data on distributed# processes first then shard on worker processestotal_workers * world_sizeglobal_worker_id global_worker_id * world_size rank_id# For BC, use default SHARDING_PRIORITIEStorch.utils.data.graph_settings.apply_sharding(datapipe, total_workers, global_worker_id)if worker_init_fn is not None:worker_init_fn(worker_id)def _share_dist_seed(generator, pg):_shared_seed torch.empty((), dtypetorch.int64).random_(generatorgenerator)if isinstance(pg, dist.ProcessGroup):dist.broadcast(_shared_seed, src0, grouppg)return _shared_seed.item()class DataLoader(Generic[_T_co]):rData loader combines a dataset and a sampler, and provides an iterable over the given dataset.The :class:~torch.utils.data.DataLoader supports both map-style anditerable-style datasets with single- or multi-process loading, customizingloading order and optional automatic batching (collation) and memory pinning.See :py:mod:torch.utils.data documentation page for more details.Args:dataset (Dataset): dataset from which to load the data.batch_size (int, optional): how many samples per batch to load(default: 1).shuffle (bool, optional): set to True to have the data reshuffledat every epoch (default: False).sampler (Sampler or Iterable, optional): defines the strategy to drawsamples from the dataset. Can be any Iterable with __len__implemented. If specified, :attr:shuffle must not be specified.batch_sampler (Sampler or Iterable, optional): like :attr:sampler, butreturns a batch of indices at a time. Mutually exclusive with:attr:batch_size, :attr:shuffle, :attr:sampler,and :attr:drop_last.num_workers (int, optional): how many subprocesses to use for dataloading. 0 means that the data will be loaded in the main process.(default: 0)collate_fn (Callable, optional): merges a list of samples to form amini-batch of Tensor(s). Used when using batched loading from amap-style dataset.pin_memory (bool, optional): If True, the data loader will copy Tensorsinto device/CUDA pinned memory before returning them. If your data elementsare a custom type, or your :attr:collate_fn returns a batch that is a custom type,see the example below.drop_last (bool, optional): set to True to drop the last incomplete batch,if the dataset size is not divisible by the batch size. If False andthe size of dataset is not divisible by the batch size, then the last batchwill be smaller. (default: False)timeout (numeric, optional): if positive, the timeout value for collecting a batchfrom workers. Should always be non-negative. (default: 0)worker_init_fn (Callable, optional): If not None, this will be called on eachworker subprocess with the worker id (an int in [0, num_workers - 1]) asinput, after seeding and before data loading. (default: None)multiprocessing_context (str or multiprocessing.context.BaseContext, optional): IfNone, the default multiprocessing context_ of your operating system willbe used. (default: None)generator (torch.Generator, optional): If not None, this RNG will be usedby RandomSampler to generate random indexes and multiprocessing to generatebase_seed for workers. (default: None)prefetch_factor (int, optional, keyword-only arg): Number of batches loadedin advance by each worker. 2 means there will be a total of2 * num_workers batches prefetched across all workers. (default value dependson the set value for num_workers. If value of num_workers0 default is None.Otherwise, if value of num_workers 0 default is 2).persistent_workers (bool, optional): If True, the data loader will not shut downthe worker processes after a dataset has been consumed once. This allows tomaintain the workers Dataset instances alive. (default: False)pin_memory_device (str, optional): the device to :attr:pin_memory to if pin_memory isTrue... warning:: If the spawn start method is used, :attr:worker_init_fncannot be an unpicklable object, e.g., a lambda function. See:ref:multiprocessing-best-practices on more details relatedto multiprocessing in PyTorch... warning:: len(dataloader) heuristic is based on the length of the sampler used.When :attr:dataset is an :class:~torch.utils.data.IterableDataset,it instead returns an estimate based on len(dataset) / batch_size, with properrounding depending on :attr:drop_last, regardless of multi-process loadingconfigurations. This represents the best guess PyTorch can make because PyTorchtrusts user :attr:dataset code in correctly handling multi-processloading to avoid duplicate data.However, if sharding results in multiple workers having incomplete last batches,this estimate can still be inaccurate, because (1) an otherwise complete batch canbe broken into multiple ones and (2) more than one batch worth of samples can bedropped when :attr:drop_last is set. Unfortunately, PyTorch can not detect suchcases in general.See Dataset Types_ for more details on these two types of datasets and how:class:~torch.utils.data.IterableDataset interacts withMulti-process data loading_... warning:: See :ref:reproducibility, and :ref:dataloader-workers-random-seed, and:ref:data-loading-randomness notes for random seed related questions... _multiprocessing context:https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methodsdataset: Dataset[_T_co]batch_size: Optional[int]num_workers: intpin_memory: booldrop_last: booltimeout: floatsampler: Union[Sampler, Iterable]pin_memory_device: strprefetch_factor: Optional[int]_iterator: Optional[_BaseDataLoaderIter]__initialized Falsedef __init__(self,dataset: Dataset[_T_co],batch_size: Optional[int] 1,shuffle: Optional[bool] None,sampler: Union[Sampler, Iterable, None] None,batch_sampler: Union[Sampler[List], Iterable[List], None] None,num_workers: int 0,collate_fn: Optional[_collate_fn_t] None,pin_memory: bool False,drop_last: bool False,timeout: float 0,worker_init_fn: Optional[_worker_init_fn_t] None,multiprocessing_contextNone,generatorNone,*,prefetch_factor: Optional[int] None,persistent_workers: bool False,pin_memory_device: str ,):torch._C._log_api_usage_once(python.data_loader)if num_workers 0:raise ValueError(num_workers option should be non-negative; use num_workers0 to disable multiprocessing.)if timeout 0:raise ValueError(timeout option should be non-negative)if num_workers 0 and prefetch_factor is not None:raise ValueError(prefetch_factor option could only be specified in multiprocessing.let num_workers 0 to enable multiprocessing, otherwise set prefetch_factor to None.)elif num_workers 0 and prefetch_factor is None:prefetch_factor 2elif prefetch_factor is not None and prefetch_factor 0:raise ValueError(prefetch_factor option should be non-negative)if persistent_workers and num_workers 0:raise ValueError(persistent_workers option needs num_workers 0)self.dataset datasetself.num_workers num_workersself.prefetch_factor prefetch_factorself.pin_memory pin_memoryself.pin_memory_device pin_memory_deviceself.timeout timeoutself.worker_init_fn worker_init_fnself.multiprocessing_context multiprocessing_context# Adds forward compatibilities so classic DataLoader can work with DataPipes:# _DataPipeSerializationWrapper container makes it easier to serialize without redefining picklerif isinstance(self.dataset, IterDataPipe):self.dataset _IterDataPipeSerializationWrapper(self.dataset)elif isinstance(self.dataset, MapDataPipe):self.dataset _MapDataPipeSerializationWrapper(self.dataset)# Arg-check dataset related before checking samplers because we want to# tell users that iterable-style datasets are incompatible with custom# samplers first, so that they dont learn that this combo doesnt work# after spending time fixing the custom sampler errors.if isinstance(dataset, IterableDataset):self._dataset_kind _DatasetKind.Iterable# NOTE [ Custom Samplers and IterableDataset ]## IterableDataset does not support custom batch_sampler or# sampler since the key is irrelevant (unless we support# generator-style dataset one day...).## For sampler, we always create a dummy sampler. This is an# infinite sampler even when the dataset may have an implemented# finite __len__ because in multi-process data loading, naive# settings will return duplicated data (which may be desired), and# thus using a sampler with length matching that of dataset will# cause data lost (you may have duplicates of the first couple# batches, but never see anything afterwards). Therefore,# Iterabledataset always uses an infinite sampler, an instance of# _InfiniteConstantSampler defined above.## A custom batch_sampler essentially only controls the batch size.# However, it is unclear how useful it would be since an iterable-style# dataset can handle that within itself. Moreover, it is pointless# in multi-process data loading as the assignment order of batches# to workers is an implementation detail so users can not control# how to batchify each workers iterable. Thus, we disable this# option. If this turns out to be useful in future, we can re-enable# this, and support custom samplers that specify the assignments to# specific workers.if isinstance(dataset, IterDataPipe):if shuffle is not None:dataset torch.utils.data.graph_settings.apply_shuffle_settings(dataset, shuffleshuffle)# We cannot check shuffle is not None here, since previously shuffleFalse was the default.elif shuffle not in {False, None}:raise ValueError(fDataLoader with IterableDataset: expected unspecified shuffle option, but got shuffle{shuffle})if sampler is not None:# See NOTE [ Custom Samplers and IterableDataset ]raise ValueError(fDataLoader with IterableDataset: expected unspecified sampler option, but got sampler{sampler})elif batch_sampler is not None:# See NOTE [ Custom Samplers and IterableDataset ]raise ValueError(DataLoader with IterableDataset: expected unspecified fbatch_sampler option, but got batch_sampler{batch_sampler})else:shuffle bool(shuffle)self._dataset_kind _DatasetKind.Mapif sampler is not None and shuffle:raise ValueError(sampler option is mutually exclusive with shuffle)if batch_sampler is not None:# auto_collation with custom batch_samplerif batch_size ! 1 or shuffle or sampler is not None or drop_last:raise ValueError(batch_sampler option is mutually exclusive with batch_size, shuffle, sampler, and drop_last)batch_size Nonedrop_last Falseelif batch_size is None:# no auto_collationif drop_last:raise ValueError(batch_sizeNone option disables auto-batching and is mutually exclusive with drop_last)if sampler is None: # give default samplersif self._dataset_kind _DatasetKind.Iterable:# See NOTE [ Custom Samplers and IterableDataset ]sampler _InfiniteConstantSampler()else: # map-styleif shuffle:sampler RandomSampler(dataset, generatorgenerator) # type: ignore[arg-type]else:sampler SequentialSampler(dataset) # type: ignore[arg-type]if batch_size is not None and batch_sampler is None:# auto_collation without custom batch_samplerbatch_sampler BatchSampler(sampler, batch_size, drop_last)self.batch_size batch_sizeself.drop_last drop_lastself.sampler samplerself.batch_sampler batch_samplerself.generator generatorif collate_fn is None:if self._auto_collation:collate_fn _utils.collate.default_collateelse:collate_fn _utils.collate.default_convertself.collate_fn collate_fnself.persistent_workers persistent_workersself.__initialized Trueself._IterableDataset_len_called (None # See NOTE [ IterableDataset and __len__ ])self._iterator Noneself.check_worker_number_rationality()torch.set_vital(Dataloader, enabled, True) # type: ignore[attr-defined]def _get_iterator(self) - _BaseDataLoaderIter:if self.num_workers 0:return _SingleProcessDataLoaderIter(self)else:self.check_worker_number_rationality()return _MultiProcessingDataLoaderIter(self)propertydef multiprocessing_context(self):return self.__multiprocessing_contextmultiprocessing_context.setterdef multiprocessing_context(self, multiprocessing_context):if multiprocessing_context is not None:if self.num_workers 0:if isinstance(multiprocessing_context, str):valid_start_methods torch.multiprocessing.get_all_start_methods()if multiprocessing_context not in valid_start_methods:raise ValueError(multiprocessing_context option fshould specify a valid start method in {valid_start_methods!r}, but got fmultiprocessing_context{multiprocessing_context!r})multiprocessing_context torch.multiprocessing.get_context(multiprocessing_context)if not isinstance(multiprocessing_context, python_multiprocessing.context.BaseContext):raise TypeError(multiprocessing_context option should be a valid context object or a string specifying the start method, but got fmultiprocessing_context{multiprocessing_context})else:raise ValueError(multiprocessing_context can only be used with multi-process loading (num_workers 0), but got fnum_workers{self.num_workers})self.__multiprocessing_context multiprocessing_contextdef __setattr__(self, attr, val):if self.__initialized and attr in (batch_size,batch_sampler,sampler,drop_last,dataset,persistent_workers,):raise ValueError(f{attr} attribute should not be set after {self.__class__.__name__} is initialized)super().__setattr__(attr, val)# We quote _BaseDataLoaderIter since it isnt defined yet and the definition cant be moved up# since _BaseDataLoaderIter references DataLoader.def __iter__(self) - _BaseDataLoaderIter:# When using a single worker the returned iterator should be# created everytime to avoid resetting its state# However, in the case of a multiple workers iterator# the iterator is only created once in the lifetime of the# DataLoader object so that workers can be reusedif self.persistent_workers and self.num_workers 0:if self._iterator is None:self._iterator self._get_iterator()else:self._iterator._reset(self)return self._iteratorelse:return self._get_iterator()propertydef _auto_collation(self):return self.batch_sampler is not Nonepropertydef _index_sampler(self):# The actual sampler used for generating indices for _DatasetFetcher# (see _utils/fetch.py) to read data at each time. This would be# .batch_sampler if in auto-collation mode, and .sampler otherwise.# We cant change .sampler and .batch_sampler attributes for BC# reasons.if self._auto_collation:return self.batch_samplerelse:return self.samplerdef __len__(self) - int:if self._dataset_kind _DatasetKind.Iterable:# NOTE [ IterableDataset and __len__ ]## For IterableDataset, __len__ could be inaccurate when one naively# does multi-processing data loading, since the samples will be duplicated.# However, no real use case should be actually using that behavior, so# it should count as a user error. We should generally trust user# code to do the proper thing (e.g., configure each replica differently# in __iter__), and give us the correct __len__ if they choose to# implement it (this will still throw if the dataset does not implement# a __len__).## To provide a further warning, we track if __len__ was called on the# DataLoader, save the returned value in self._len_called, and warn# if the iterator ends up yielding more than this number of samples.# Cannot statically verify that dataset is Sizedlength self._IterableDataset_len_called len(self.dataset) # type: ignore[assignment, arg-type]if (self.batch_size is not None): # IterableDataset doesnt allow custom sampler or batch_samplerfrom math import ceilif self.drop_last:length length // self.batch_sizeelse:length ceil(length / self.batch_size)return lengthelse:return len(self._index_sampler)def check_worker_number_rationality(self):# This function check whether the dataloaders worker number is rational based on# current systems resource. Current rule is that if the number of workers this# Dataloader will create is bigger than the number of logical cpus that is allowed to# use, than we will pop up a warning to let user pay attention.## eg. If current system has 2 physical CPUs with 16 cores each. And each core support 2# threads, then the total logical cpus here is 2 * 16 * 2 64. Lets say current# DataLoader process can use half of them which is 32, then the rational max number of# worker that initiated from this process is 32.# Now, lets say the created DataLoader has num_works 40, which is bigger than 32.# So the warning message is triggered to notify the user to lower the worker number if# necessary.### [Note] Please note that this function repects cpuset only when os.sched_getaffinity is# available (available in most of Linux system, but not OSX and Windows).# When os.sched_getaffinity is not available, os.cpu_count() is called instead, but# it doesnt repect cpuset.# We dont take threading into account since each worker process is single threaded# at this time.## We dont set any threading flags (eg. OMP_NUM_THREADS, MKL_NUM_THREADS, etc)# other than torch.set_num_threads to 1 in the worker process, if the passing# in functions use 3rd party modules that rely on those threading flags to determine# how many thread to create (eg. numpy, etc), then it is callers responsibility to# set those flags correctly.def _create_warning_msg(num_worker_suggest, num_worker_created, cpuset_checked):suggested_max_worker_msg (((Our suggested max number of worker in current system is {}{}, which is smaller than what this DataLoader is going to create.).format(num_worker_suggest,(if cpuset_checkedelse (cpuset is not taken into account)),))if num_worker_suggest is not Noneelse (DataLoader is not able to compute a suggested max number of worker in current system.))warn_msg (fThis DataLoader will create {num_worker_created} worker processes in total. {suggested_max_worker_msg} Please be aware that excessive worker creation might get DataLoader running slow or even freeze, lower the worker number to avoid potential slowness/freeze if necessary.)return warn_msgif not self.num_workers or self.num_workers 0:return# try to compute a suggested max number of worker based on systems resourcemax_num_worker_suggest Nonecpuset_checked Falseif hasattr(os, sched_getaffinity):try:max_num_worker_suggest len(os.sched_getaffinity(0))cpuset_checked Trueexcept Exception:passif max_num_worker_suggest is None:# os.cpu_count() could return Optional[int]# get cpu count first and check None in order to satisfy mypy checkcpu_count os.cpu_count()if cpu_count is not None:max_num_worker_suggest cpu_countif max_num_worker_suggest is None:warnings.warn(_create_warning_msg(max_num_worker_suggest, self.num_workers, cpuset_checked))returnif self.num_workers max_num_worker_suggest:warnings.warn(_create_warning_msg(max_num_worker_suggest, self.num_workers, cpuset_checked))class _BaseDataLoaderIter:def __init__(self, loader: DataLoader) - None:self._dataset loader.datasetself._shared_seed Noneself._pg Noneif isinstance(self._dataset, IterDataPipe):if dist.is_available() and dist.is_initialized():self._pg dist.new_group(backendgloo)self._shared_seed _share_dist_seed(loader.generator, self._pg)shared_rng torch.Generator()shared_rng.manual_seed(self._shared_seed)self._dataset torch.utils.data.graph_settings.apply_random_seed(self._dataset, shared_rng)self._dataset_kind loader._dataset_kindself._IterableDataset_len_called loader._IterableDataset_len_calledself._auto_collation loader._auto_collationself._drop_last loader.drop_lastself._index_sampler loader._index_samplerself._num_workers loader.num_workersws, rank _get_distributed_settings()self._world_size wsself._rank rank# for other backends, pin_memory_device need to set. if not set# default behaviour is CUDA device. if pin_memory_device is selected# and pin_memory is not set, the default behaviour false.if len(loader.pin_memory_device) 0:self._pin_memory loader.pin_memory and torch.cuda.is_available()self._pin_memory_device Noneelse:if not loader.pin_memory:warn_msg (pin memory device is set and pin_memory flag is not used then device pinned memory wont be usedplease set pin_memory to true, if you need to use the device pin memory)warnings.warn(warn_msg)self._pin_memory loader.pin_memoryself._pin_memory_device loader.pin_memory_deviceself._timeout loader.timeoutself._collate_fn loader.collate_fnself._sampler_iter iter(self._index_sampler)self._base_seed (torch.empty((), dtypetorch.int64).random_(generatorloader.generator).item())self._persistent_workers loader.persistent_workersself._num_yielded 0self._profile_name fenumerate(DataLoader)#{self.__class__.__name__}.__next__def __iter__(self) - _BaseDataLoaderIter:return selfdef _reset(self, loader, first_iterFalse):self._sampler_iter iter(self._index_sampler)self._num_yielded 0self._IterableDataset_len_called loader._IterableDataset_len_calledif isinstance(self._dataset, IterDataPipe):self._shared_seed _share_dist_seed(loader.generator, self._pg)shared_rng torch.Generator()shared_rng.manual_seed(self._shared_seed)self._dataset torch.utils.data.graph_settings.apply_random_seed(self._dataset, shared_rng)def _next_index(self):return next(self._sampler_iter) # may raise StopIterationdef _next_data(self):raise NotImplementedErrordef __next__(self) - Any:with torch.autograd.profiler.record_function(self._profile_name):if self._sampler_iter is None:# TODO(https://github.com/pytorch/pytorch/issues/76750)self._reset() # type: ignore[call-arg]data self._next_data()self._num_yielded 1if (self._dataset_kind _DatasetKind.Iterableand self._IterableDataset_len_called is not Noneand self._num_yielded self._IterableDataset_len_called):warn_msg (fLength of IterableDataset {self._dataset} was reported to be {self._IterableDataset_len_called}f(when accessing len(dataloader)), but {self._num_yielded} samples have been fetched. )if self._num_workers 0:warn_msg (For multiprocessing data-loading, this could be caused by not properly configuring the IterableDataset replica at each worker. Please see https://pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset for examples.)warnings.warn(warn_msg)return datadef __len__(self) - int:return len(self._index_sampler)def __getstate__(self):# TODO: add limited pickling support for sharing an iterator# across multiple threads for HOGWILD.# Probably the best way to do this is by moving the sample pushing# to a separate thread and then just sharing the data queue# but signalling the end is tricky without a non-blocking APIraise NotImplementedError({} cannot be pickled, self.__class__.__name__)class _SingleProcessDataLoaderIter(_BaseDataLoaderIter):def __init__(self, loader):super().__init__(loader)assert self._timeout 0assert self._num_workers 0# Adds forward compatibilities so classic DataLoader can work with DataPipes:# Taking care of distributed shardingif isinstance(self._dataset, (IterDataPipe, MapDataPipe)):# For BC, use default SHARDING_PRIORITIEStorch.utils.data.graph_settings.apply_sharding(self._dataset, self._world_size, self._rank)self._dataset_fetcher _DatasetKind.create_fetcher(self._dataset_kind,self._dataset,self._auto_collation,self._collate_fn,self._drop_last,)def _next_data(self):index self._next_index() # may raise StopIterationdata self._dataset_fetcher.fetch(index) # may raise StopIterationif self._pin_memory:data _utils.pin_memory.pin_memory(data, self._pin_memory_device)return dataclass _MultiProcessingDataLoaderIter(_BaseDataLoaderIter):rIterates once over the DataLoaders dataset, as specified by the sampler.# NOTE [ Data Loader Multiprocessing Shutdown Logic ]## Preliminary:## Our data model looks like this (queues are indicated with curly brackets):## main process ||# | ||# {index_queue} ||# | ||# worker processes || DATA# | ||# {worker_result_queue} || FLOW# | ||# pin_memory_thread of main process || DIRECTION# | ||# {data_queue} ||# | ||# data output \/## P.S. worker_result_queue and pin_memory_thread part may be omitted if# pin_memoryFalse.### Terminating multiprocessing logic requires very careful design. In# particular, we need to make sure that## 1. The iterator gracefully exits the workers when its last reference is# gone or it is depleted.## In this case, the workers should be gracefully exited because the# main process may still need to continue to run, and we want cleaning# up code in the workers to be executed (e.g., releasing GPU memory).# Naturally, we implement the shutdown logic in __del__ of# DataLoaderIterator.## We delay the discussion on the logic in this case until later.## 2. The iterator exits the workers when the loader process and/or worker# processes exits normally or with error.## We set all workers and pin_memory_thread to have daemonTrue.## You may ask, why cant we make the workers non-daemonic, and# gracefully exit using the same logic as we have in __del__ when the# iterator gets deleted (see 1 above)?## First of all, __del__ is **not** guaranteed to be called when# interpreter exits. Even if it is called, by the time it executes,# many Python core library resources may already be freed, and even# simple things like acquiring an internal lock of a queue may hang.# Therefore, in this case, we actually need to prevent __del__ from# being executed, and rely on the automatic termination of daemonic# children.## Thus, we register an atexit hook that sets a global flag# _utils.python_exit_status. Since atexit hooks are executed in the# reverse order of registration, we are guaranteed that this flag is# set before library resources we use are freed (which, at least in# CPython, is done via an atexit handler defined in# multiprocessing/util.py# https://github.com/python/cpython/blob/c606624af8d4cb3b4a052fb263bb983b3f87585b/Lib/multiprocessing/util.py#L320-L362# registered when an object requiring this mechanism is first# created, e.g., mp.Queue# https://github.com/python/cpython/blob/c606624af8d4cb3b4a052fb263bb983b3f87585b/Lib/multiprocessing/context.py#L100-L103# https://github.com/python/cpython/blob/c606624af8d4cb3b4a052fb263bb983b3f87585b/Lib/multiprocessing/queues.py#L29# )## So in __del__, we check if _utils.python_exit_status is set or# None (freed), and perform no-op if so.## However, simply letting library clean-up codes run can also be bad,# because such codes (i.e., multiprocessing.util._exit_function())# include join putting threads for mp.Queue, which can be blocking.# Hence, the main process putting threads are called with# cancel_join_thread at creation. See later section# [ 3b. A process wont hang when putting into a queue; ]# for more details.## Here are two example cases where library clean-up codes can run# before __del__ is called:## 1. If we hold onto a reference to the iterator, it more often# than not tries to do multiprocessing library cleaning before# clearing the alive referenced objects (https://github.com/pytorch/pytorch/issues/48666)# and thus prevents our cleaning-up code to run first.## 2. A similar issue araises when a DataLoader is used in a subprocess.# When a process ends, it shuts the all its daemonic children# down with a SIGTERM (instead of joining them without a timeout).# Simiarly for threads, but by a different mechanism. This fact,# together with a few implementation details of multiprocessing, forces# us to make workers daemonic. All of our problems arise when a# DataLoader is used in a subprocess, and are caused by multiprocessing# code which looks more or less like this:## try:# your_function_using_a_dataloader()# finally:# multiprocessing.util._exit_function()## The joining/termination mentioned above happens inside# _exit_function(). Now, if your_function_using_a_dataloader()# throws, the stack trace stored in the exception will prevent the# frame which uses DataLoaderIter to be freed. If the frame has any# reference to the DataLoaderIter (e.g., in a method of the iter),# its __del__, which starts the shutdown procedure, will not be# called. That, in turn, means that workers arent notified. Attempting# to join in _exit_function will then result in a hang.## For context, _exit_function is also registered as an atexit call.# So it is unclear to me (ssnl) why this is needed in a finally block.# The code dates back to 2008 and there is no comment on the original# PEP 371 or patch https://bugs.python.org/issue3050 (containing both# the finally block and the atexit registration) that explains this.### Finally, another choice is to just shutdown workers with logic in 1# above whenever we see an error in next. This isnt ideal because# a. It prevents users from using try-catch to resume data loading.# b. It doesnt prevent hanging if users have references to the# iterator.## 3. All processes exit if any of them die unexpectedly by fatal signals.## As shown above, the workers are set as daemonic children of the main# process. However, automatic cleaning-up of such child processes only# happens if the parent process exits gracefully (e.g., not via fatal# signals like SIGKILL). So we must ensure that each process will exit# even the process that should send/receive data to/from it were# killed, i.e.,## a. A process wont hang when getting from a queue.## Even with carefully designed data dependencies (i.e., a put()# always corresponding to a get()), hanging on get() can still# happen when data in queue is corrupted (e.g., due to# cancel_join_thread or unexpected exit).## For child exit, we set a timeout whenever we try to get data# from data_queue, and check the workers status on each timeout# and error.# See _DataLoaderiter._get_batch() and# _DataLoaderiter._try_get_data() for details.## Additionally, for child exit on non-Windows platforms, we also# register a SIGCHLD handler (which is supported on Windows) on# the main process, which checks if any of the workers fail in the# (Python) handler. This is more efficient and faster in detecting# worker failures, compared to only using the above mechanism.# See DataLoader.cpp and _utils/signal_handling.py for details.## For .get() calls where the sender(s) is not the workers, we# guard them with timeouts, and check the status of the sender# when timeout happens:# in the workers, the _utils.worker.ManagerWatchdog class# checks the status of the main process.# if pin_memoryTrue, when getting from pin_memory_thread,# check pin_memory_thread status periodically until .get()# returns or see that pin_memory_thread died.## b. A process wont hang when putting into a queue;## We use mp.Queue which has a separate background thread to put# objects from an unbounded buffer array. The background thread is# daemonic and usually automatically joined when the process# *exits*.## In case that the receiver has ended abruptly while# reading from the pipe, the join will hang forever. The usual# solution for this in Python is calling q.cancel_join_thread,# which prevents automatically joining it when finalizing# (exiting).## Nonetheless, cancel_join_thread must only be called when the# queue is **not** going to be read from or write into by another# process, because it may hold onto a lock or leave corrupted data# in the queue, leading other readers/writers to hang.## Hence,# For worker processes, we only do so (for their output# queues, i.e., worker_result_queue) before exiting.# For pin_memory_thread, its output queue data_queue is a# queue.Queue that does blocking put if the queue is full.# So there is no above problem, but as a result, in# _pin_memory_loop, we do need to wrap the put in a loop# that breaks not only upon success, but also when the main# process stops reading, i.e., is shutting down.# For loader process, we cancel_join_thread() for all# _index_queues because the whole purpose of workers and# pin_memory_thread is to serve the loader process. If# loader process is already exiting, we dont really care if# the queues are corrupted.### Now lets get back to 1:# how we gracefully exit the workers when the last reference to the# iterator is gone.## To achieve this, we implement the following logic along with the design# choices mentioned above:## workers_done_event:# A multiprocessing.Event shared among the main process and all worker# processes. This is used to signal the workers that the iterator is# shutting down. After it is set, they will not send processed data to# queues anymore, and only wait for the final None before exiting.# done_event isnt strictly needed. I.e., we can just check for None# from the input queue, but it allows us to skip wasting resources# processing data if we are already shutting down.## pin_memory_thread_done_event:# A threading.Event for a similar purpose to that of# workers_done_event, but is for the pin_memory_thread. The reason# that separate events are needed is that pin_memory_thread reads from# the output queue of the workers. But the workers, upon seeing that# workers_done_event is set, only wants to see the final None, and is# not required to flush all data in the output queue (e.g., it may call# cancel_join_thread on that queue if its IterableDataset iterator# happens to exhaust coincidentally, which is out of the control of the# main process). Thus, since we will exit pin_memory_thread before the# workers (see below), two separete events are used.## NOTE: In short, the protocol is that the main process will set these# done_events and then the corresponding processes/threads a None,# and that they may exit at any time after receiving the None.## NOTE: Using None as the final signal is valid, since normal data will# always be a 2-tuple with the 1st element being the index of the data# transferred (different from dataset index/key), and the 2nd being# either the dataset key or the data sample (depending on which part# of the data model the queue is at).## [ worker processes ]# While loader process is alive:# Get from index_queue.# If get anything else,# Check workers_done_event.# If set, continue to next iteration# i.e., keep getting until see the None, then exit.# Otherwise, process data:# If is fetching from an IterableDataset and the iterator# is exhausted, send an _IterableDatasetStopIteration# object to signal iteration end. The main process, upon# receiving such an object, will send None to this# worker and not use the corresponding index_queue# anymore.# If timed out,# No matter workers_done_event is set (still need to see None)# or not, must continue to next iteration.# (outside loop)# If workers_done_event is set, (this can be False with IterableDataset)# data_queue.cancel_join_thread(). (Everything is ending here:# main process wont read from it;# other workers will also call# cancel_join_thread.)## [ pin_memory_thread ]# # No need to check main thread. If this thread is alive, the main loader# # thread must be alive, because this thread is set as daemonic.# While pin_memory_thread_done_event is not set:# Get from worker_result_queue.# If timed out, continue to get in the next iteration.# Otherwise, process data.# While pin_memory_thread_done_event is not set:# Put processed data to data_queue (a queue.Queue with blocking put)# If timed out, continue to put in the next iteration.# Otherwise, break, i.e., continuing to the out loop.## NOTE: we dont check the status of the main thread because# 1. if the process is killed by fatal signal, pin_memory_thread# ends.# 2. in other cases, either the cleaning-up in __del__ or the# automatic exit of daemonic thread will take care of it.# This wont busy-wait either because .get(timeout) does not# busy-wait.## [ main process ]# In the DataLoader Iters __del__# b. Exit pin_memory_thread# i. Set pin_memory_thread_done_event.# ii Put None in worker_result_queue.# iii. Join the pin_memory_thread.# iv. worker_result_queue.cancel_join_thread().## c. Exit the workers.# i. Set workers_done_event.# ii. Put None in each workers index_queue.# iii. Join the workers.# iv. Call .cancel_join_thread() on each workers index_queue.## NOTE: (c) is better placed after (b) because it may leave corrupted# data in worker_result_queue, which pin_memory_thread# reads from, in which case the pin_memory_thread can only# happen at timing out, which is slow. Nonetheless, same thing# happens if a worker is killed by signal at unfortunate times,# but in other cases, we are better off having a non-corrupted# worker_result_queue for pin_memory_thread.## NOTE: If pin_memoryFalse, there is no pin_memory_thread and (b)# can be omitted## NB: done_events isnt strictly needed. E.g., we can just check for# None from index_queue, but it allows us to skip wasting resources# processing indices already in index_queue if we are already shutting# down.def __init__(self, loader):super().__init__(loader)self._prefetch_factor loader.prefetch_factorassert self._num_workers 0assert self._prefetch_factor 0if loader.multiprocessing_context is None:multiprocessing_context torch.multiprocessingelse:multiprocessing_context loader.multiprocessing_contextself._worker_init_fn loader.worker_init_fn# Adds forward compatibilities so classic DataLoader can work with DataPipes:# Additional worker init function will take care of sharding in MP and Distributedif isinstance(self._dataset, (IterDataPipe, MapDataPipe)):self._worker_init_fn functools.partial(_sharding_worker_init_fn,self._worker_init_fn,self._world_size,self._rank,)# No certainty which module multiprocessing_context isself._worker_result_queue multiprocessing_context.Queue() # type: ignore[var-annotated]self._worker_pids_set Falseself._shutdown Falseself._workers_done_event multiprocessing_context.Event()self._index_queues []self._workers []for i in range(self._num_workers):# No certainty which module multiprocessing_context isindex_queue multiprocessing_context.Queue() # type: ignore[var-annotated]# Need to cancel_join_thread here!# See sections (2) and (3b) above.index_queue.cancel_join_thread()w multiprocessing_context.Process(target_utils.worker._worker_loop,args(self._dataset_kind,self._dataset,index_queue,self._worker_result_queue,self._workers_done_event,self._auto_collation,self._collate_fn,self._drop_last,self._base_seed,self._worker_init_fn,i,self._num_workers,self._persistent_workers,self._shared_seed,),)w.daemon True# NB: Process.start() actually take some time as it needs to# start a process and pass the arguments over via a pipe.# Therefore, we only add a worker to self._workers list after# it started, so that we do not call .join() if program dies# before it starts, and __del__ tries to join but will get:# AssertionError: can only join a started process.w.start()self._index_queues.append(index_queue)self._workers.append(w)if self._pin_memory:self._pin_memory_thread_done_event threading.Event()# Queue is not type-annotatedself._data_queue queue.Queue() # type: ignore[var-annotated]if self._pin_memory_device xpu:current_device torch.xpu.current_device() # type: ignore[attr-defined]elif self._pin_memory_device torch._C._get_privateuse1_backend_name():custom_device_mod getattr(torch, torch._C._get_privateuse1_backend_name())current_device custom_device_mod.current_device()else:current_device torch.cuda.current_device() # choose cuda for defaultpin_memory_thread threading.Thread(target_utils.pin_memory._pin_memory_loop,args(self._worker_result_queue,self._data_queue,current_device,self._pin_memory_thread_done_event,self._pin_memory_device,),)pin_memory_thread.daemon Truepin_memory_thread.start()# Similar to workers (see comment above), we only register# pin_memory_thread once it is started.self._pin_memory_thread pin_memory_threadelse:self._data_queue self._worker_result_queue # type: ignore[assignment]# In some rare cases, persistent workers (daemonic processes)# would be terminated before __del__ of iterator is invoked# when main process exits# It would cause failure when pin_memory_thread tries to read# corrupted data from worker_result_queue# atexit is used to shutdown thread and child processes in the# right sequence before main process exitsif self._persistent_workers and self._pin_memory:import atexitfor w in self._workers:atexit.register(_MultiProcessingDataLoaderIter._clean_up_worker, w)# .pid can be None only before process is spawned (not the case, so ignore)_utils.signal_handling._set_worker_pids(id(self), tuple(w.pid for w in self._workers)) # type: ignore[misc]_utils.signal_handling._set_SIGCHLD_handler()self._worker_pids_set Trueself._reset(loader, first_iterTrue)def _reset(self, loader, first_iterFalse):super()._reset(loader, first_iter)self._send_idx 0 # idx of the next task to be sent to workersself._rcvd_idx 0 # idx of the next task to be returned in __next__# information about data not yet yielded, i.e., tasks w/ indices in range [rcvd_idx, send_idx).# map: task idx - (worker_id,) if data isnt fetched (outstanding)# \ (worker_id, data) if data is already fetched (out-of-order)self._task_info {}self._tasks_outstanding (0 # always equal to count(v for v in task_info.values() if len(v) 1))# A list of booleans representing whether each worker still has work to# do, i.e., not having exhausted its iterable dataset object. It always# contains all Trues if not using an iterable-style dataset# (i.e., if kind ! Iterable).# Not that this indicates that a worker still has work to do *for this epoch*.# It does not mean that a worker is dead. In case of _persistent_workers,# the worker will be reset to available in the next epoch.self._workers_status [True for i in range(self._num_workers)]# Reset the worker queue cycle so it resumes next epoch at worker 0self._worker_queue_idx_cycle itertools.cycle(range(self._num_workers))# We resume the prefetching in case it was enabledif not first_iter:for idx in range(self._num_workers):self._index_queues[idx].put(_utils.worker._ResumeIteration(self._shared_seed))resume_iteration_cnt self._num_workerswhile resume_iteration_cnt 0:return_idx, return_data self._get_data()if isinstance(return_idx, _utils.worker._ResumeIteration):assert return_data is Noneresume_iteration_cnt - 1# prime the prefetch loopfor _ in range(self._prefetch_factor * self._num_workers):self._try_put_index()def _try_get_data(self, timeout_utils.MP_STATUS_CHECK_INTERVAL):# Tries to fetch data from self._data_queue once for a given timeout.# This can also be used as inner loop of fetching without timeout, with# the sender status as the loop condition.## This raises a RuntimeError if any worker died expectedly. This error# can come from either the SIGCHLD handler in _utils/signal_handling.py# (only for non-Windows platforms), or the manual check below on errors# and timeouts.## Returns a 2-tuple:# (bool: whether successfully get data, any: data if successful else None)try:data self._data_queue.get(timeouttimeout)return (True, data)except Exception as e:# At timeout and error, we manually check whether any worker has# failed. Note that this is the only mechanism for Windows to detect# worker failures.failed_workers []for worker_id, w in enumerate(self._workers):if self._workers_status[worker_id] and not w.is_alive():failed_workers.append(w)self._mark_worker_as_unavailable(worker_id)if len(failed_workers) 0:pids_str , .join(str(w.pid) for w in failed_workers)raise RuntimeError(fDataLoader worker (pid(s) {pids_str}) exited unexpectedly) from eif isinstance(e, queue.Empty):return (False, None)import errnoimport tempfiletry:# Raise an exception if we are this close to the FDs limit.# Apparently, trying to open only one file is not a sufficient# test.# See NOTE [ DataLoader on Linux and open files limit ]fds_limit_margin 10fs [tempfile.NamedTemporaryFile() for i in range(fds_limit_margin)]except OSError as e:if e.errno errno.EMFILE:raise RuntimeError(Too many open files. Communication with the workers is no longer possible. Please increase the limit using ulimit -n in the shell or change the sharing strategy by calling torch.multiprocessing.set_sharing_strategy(file_system) at the beginning of your code) from Noneraise# NOTE [ DataLoader on Linux and open files limit ]## On Linux when DataLoader is used with multiprocessing we pass the data between# the root process and the workers through SHM files. We remove those files from# the filesystem as soon as they are created and keep them alive by# passing around their file descriptors through AF_UNIX sockets. (See# docs/source/multiprocessing.rst and Multiprocessing Technical Notes in# the wiki (https://github.com/pytorch/pytorch/wiki).)## This sometimes leads us to exceeding the open files limit. When that happens,# and the offending file descriptor is coming over a socket, the socket Python# package silently strips the file descriptor from the message, setting only the# MSG_CTRUNC flag (which might be a bit misleading since the manpage says that# it _indicates that some control data were discarded due to lack of space in# the buffer for ancillary data_). This might reflect the C implementation of# AF_UNIX sockets.## This behaviour can be reproduced with the script and instructions at the# bottom of this note.## When that happens, the standard Python multiprocessing (and not# torch.multiprocessing) raises a RuntimeError: received 0 items of ancdata## Sometimes, instead of the FD being stripped, you may get an OSError:# Too many open files, both in the script below and in DataLoader. However,# this is rare and seems to be nondeterministic.### #!/usr/bin/env python3# import sys# import socket# import os# import array# import shutil# import socket### if len(sys.argv) ! 4:# print(Usage: , sys.argv[0], tmp_dirname iteration (send|recv))# sys.exit(1)## if __name__ __main__:# dirname sys.argv[1]# sock_path dirname /sock# iterations int(sys.argv[2])# def dummy_path(i):# return dirname / str(i) .dummy### if sys.argv[3] send:# while not os.path.exists(sock_path):# pass# client socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)# client.connect(sock_path)# for i in range(iterations):# fd os.open(dummy_path(i), os.O_WRONLY | os.O_CREAT)# ancdata array.array(i, [fd])# msg bytes([i % 256])# print(Sending fd , fd, (iteration #, i, ))# client.sendmsg([msg], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, ancdata)])### else:# assert sys.argv[3] recv## if os.path.exists(dirname):# raise Exception(Directory exists)## os.mkdir(dirname)## print(Opening socket...)# server socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)# server.bind(sock_path)## print(Listening...)# for i in range(iterations):# a array.array(i)# msg, ancdata, flags, addr server.recvmsg(1, socket.CMSG_SPACE(a.itemsize))# assert(len(ancdata) 1)# cmsg_level, cmsg_type, cmsg_data ancdata[0]# a.frombytes(cmsg_data)# print(Received fd , a[0], (iteration #, i, ))## shutil.rmtree(dirname)## Steps to reproduce:## 1. Run two shells and set lower file descriptor limit in the receiving one:# (shell1) ulimit -n 1020# (shell2) ulimit -n 1022## 2. Run the script above with the recv option in the first shell# (shell1) ./test_socket.py sock_tmp 1017 recv## 3. Run the script with the send option in the second shell:# (shell2) ./test_socket.py sock_tmp 1017 senddef _get_data(self):# Fetches data from self._data_queue.## We check workers status every MP_STATUS_CHECK_INTERVAL seconds,# which we achieve by running self._try_get_data(timeoutMP_STATUS_CHECK_INTERVAL)# in a loop. This is the only mechanism to detect worker failures for# Windows. For other platforms, a SIGCHLD handler is also used for# worker failure detection.## If pin_memoryTrue, we also need check if pin_memory_thread had# died at timeouts.if self._timeout 0:success, data self._try_get_data(self._timeout)if success:return dataelse:raise RuntimeError(fDataLoader timed out after {self._timeout} seconds)elif self._pin_memory:while self._pin_memory_thread.is_alive():success, data self._try_get_data()if success:return dataelse:# while condition is false, i.e., pin_memory_thread died.raise RuntimeError(Pin memory thread exited unexpectedly)# In this case, self._data_queue is a queue.Queue,. But we dont# need to call .task_done() because we dont use .join().else:while True:success, data self._try_get_data()if success:return datadef _next_data(self):while True:# If the worker responsible for self._rcvd_idx has already ended# and was unable to fulfill this task (due to exhausting an IterableDataset),# we try to advance self._rcvd_idx to find the next valid index.## This part needs to run in the loop because both the self._get_data()# call and _IterableDatasetStopIteration check below can mark# extra worker(s) as dead.while self._rcvd_idx self._send_idx:info self._task_info[self._rcvd_idx]worker_id info[0]if (len(info) 2 or self._workers_status[worker_id]): # has data or is still activebreakdel self._task_info[self._rcvd_idx]self._rcvd_idx 1else:# no valid self._rcvd_idx is found (i.e., didnt break)if not self._persistent_workers:self._shutdown_workers()raise StopIteration# Now self._rcvd_idx is the batch index we want to fetch# Check if the next sample has already been generatedif len(self._task_info[self._rcvd_idx]) 2:data self._task_info.pop(self._rcvd_idx)[1]return self._process_data(data)assert not self._shutdown and self._tasks_outstanding 0idx, data self._get_data()self._tasks_outstanding - 1if self._dataset_kind _DatasetKind.Iterable:# Check for _IterableDatasetStopIterationif isinstance(data, _utils.worker._IterableDatasetStopIteration):if self._persistent_workers:self._workers_status[data.worker_id] Falseelse:self._mark_worker_as_unavailable(data.worker_id)self._try_put_index()continueif idx ! self._rcvd_idx:# store out-of-order samplesself._task_info[idx] (data,)else:del self._task_info[idx]return self._process_data(data)def _try_put_index(self):assert self._tasks_outstanding self._prefetch_factor * self._num_workerstry:index self._next_index()except StopIteration:returnfor _ in range(self._num_workers): # find the next active worker, if anyworker_queue_idx next(self._worker_queue_idx_cycle)if self._workers_status[worker_queue_idx]:breakelse:# not found (i.e., didnt break)returnself._index_queues[worker_queue_idx].put((self._send_idx, index)) # type: ignore[possibly-undefined]self._task_info[self._send_idx] (worker_queue_idx,)self._tasks_outstanding 1self._send_idx 1def _process_data(self, data):self._rcvd_idx 1self._try_put_index()if isinstance(data, ExceptionWrapper):data.reraise()return datadef _mark_worker_as_unavailable(self, worker_id, shutdownFalse):# Mark a worker as having finished its work e.g., due to# exhausting an IterableDataset. This should be used only when this# _MultiProcessingDataLoaderIter is going to continue running.assert self._workers_status[worker_id] or (self._persistent_workers and shutdown)# Signal termination to that specific worker.q self._index_queues[worker_id]# Indicate that no more data will be put on this queue by the current# process.q.put(None)# Note that we dont actually join the worker here, nor do we remove the# workers pid from C side struct because (1) joining may be slow, and# (2) since we dont join, the worker may still raise error, and we# prefer capturing those, rather than ignoring them, even though they# are raised after the worker has finished its job.# Joinning is deferred to _shutdown_workers, which it is called when# all workers finish their jobs (e.g., IterableDataset replicas) or# when this iterator is garbage collected.self._workers_status[worker_id] Falseassert self._workers_done_event.is_set() shutdowndef _shutdown_workers(self):# Called when shutting down this _MultiProcessingDataLoaderIter.# See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for details on# the logic of this function.if (_utils is Noneor _utils.python_exit_status is Trueor _utils.python_exit_status is None):# See (2) of the note. If Python is shutting down, do no-op.return# Normal exit when last reference is gone / iterator is depleted.# See (1) and the second half of the note.if not self._shutdown:self._shutdown Truetry:# Normal exit when last reference is gone / iterator is depleted.# See (1) and the second half of the note.# Exit pin_memory_thread first because exiting workers may leave# corrupted data in worker_result_queue which pin_memory_thread# reads from.if hasattr(self, _pin_memory_thread):# Use hasattr in case error happens before we set the attribute.self._pin_memory_thread_done_event.set()# Send something to pin_memory_thread in case it is waiting# so that it can wake up and check pin_memory_thread_done_eventself._worker_result_queue.put((None, None))self._pin_memory_thread.join()self._worker_result_queue.cancel_join_thread()self._worker_result_queue.close()# Exit workers now.self._workers_done_event.set()for worker_id in range(len(self._workers)):# Get number of workers from len(self._workers) instead of# self._num_workers in case we error before starting all# workers.# If we are using workers_status with persistent_workers# we have to shut it down because the worker is pausedif self._persistent_workers or self._workers_status[worker_id]:self._mark_worker_as_unavailable(worker_id, shutdownTrue)for w in self._workers:# We should be able to join here, but in case anything went# wrong, we set a timeout and if the workers fail to join,# they are killed in the finally block.w.join(timeout_utils.MP_STATUS_CHECK_INTERVAL)for q in self._index_queues:q.cancel_join_thread()q.close()finally:# Even though all this function does is putting into queues that# we have called cancel_join_thread on, weird things can# happen when a worker is killed by a signal, e.g., hanging in# Event.set(). So we need to guard this with SIGCHLD handler,# and remove pids from the C side data structure only at the# end.## FIXME: Unfortunately, for Windows, we are missing a worker# error detection mechanism here in this function, as it# doesnt provide a SIGCHLD handler.if self._worker_pids_set:_utils.signal_handling._remove_worker_pids(id(self))self._worker_pids_set Falsefor w in self._workers:if w.is_alive():# Existing mechanisms try to make the workers exit# peacefully, but in case that we unfortunately reach# here, which we shouldnt, (e.g., pytorch/pytorch#39570),# we kill the worker.w.terminate()# staticmethod is used to remove reference to _MultiProcessingDataLoaderIterstaticmethoddef _clean_up_worker(w):try:w.join(timeout_utils.MP_STATUS_CHECK_INTERVAL)finally:if w.is_alive():w.terminate()def __del__(self):self._shutdown_workers()
文章转载自： http://www.morning.bqdpy.cn.gov.cn.bqdpy.cn http://www.morning.tnthd.cn.gov.cn.tnthd.cn http://www.morning.rknhd.cn.gov.cn.rknhd.cn http://www.morning.bkkgt.cn.gov.cn.bkkgt.cn http://www.morning.wsxxq.cn.gov.cn.wsxxq.cn http://www.morning.qjxkx.cn.gov.cn.qjxkx.cn http://www.morning.qkdbz.cn.gov.cn.qkdbz.cn http://www.morning.qxwgx.cn.gov.cn.qxwgx.cn http://www.morning.njpny.cn.gov.cn.njpny.cn http://www.morning.wgzgr.cn.gov.cn.wgzgr.cn http://www.morning.nyjgm.cn.gov.cn.nyjgm.cn http://www.morning.nkdmd.cn.gov.cn.nkdmd.cn http://www.morning.lbbrw.cn.gov.cn.lbbrw.cn http://www.morning.qdxwf.cn.gov.cn.qdxwf.cn http://www.morning.hwlk.cn.gov.cn.hwlk.cn http://www.morning.sjwws.cn.gov.cn.sjwws.cn http://www.morning.mlntx.cn.gov.cn.mlntx.cn http://www.morning.inheatherskitchen.com.gov.cn.inheatherskitchen.com http://www.morning.nspbj.cn.gov.cn.nspbj.cn http://www.morning.rui931.cn.gov.cn.rui931.cn http://www.morning.jfnlj.cn.gov.cn.jfnlj.cn http://www.morning.bqts.cn.gov.cn.bqts.cn http://www.morning.zgdnz.cn.gov.cn.zgdnz.cn http://www.morning.jhxtm.cn.gov.cn.jhxtm.cn http://www.morning.tqsnd.cn.gov.cn.tqsnd.cn http://www.morning.mswkd.cn.gov.cn.mswkd.cn http://www.morning.xzrbd.cn.gov.cn.xzrbd.cn http://www.morning.hpxxq.cn.gov.cn.hpxxq.cn http://www.morning.jkcpl.cn.gov.cn.jkcpl.cn http://www.morning.wqkzf.cn.gov.cn.wqkzf.cn http://www.morning.dpruuode.cn.gov.cn.dpruuode.cn http://www.morning.kqxng.cn.gov.cn.kqxng.cn http://www.morning.rtspr.cn.gov.cn.rtspr.cn http://www.morning.pqnpd.cn.gov.cn.pqnpd.cn http://www.morning.jbtwq.cn.gov.cn.jbtwq.cn http://www.morning.rhzzf.cn.gov.cn.rhzzf.cn http://www.morning.qxgmp.cn.gov.cn.qxgmp.cn http://www.morning.thpns.cn.gov.cn.thpns.cn http://www.morning.jpfpc.cn.gov.cn.jpfpc.cn http://www.morning.glswq.cn.gov.cn.glswq.cn http://www.morning.qhln.cn.gov.cn.qhln.cn http://www.morning.gccdr.cn.gov.cn.gccdr.cn http://www.morning.wsxxq.cn.gov.cn.wsxxq.cn http://www.morning.dighk.com.gov.cn.dighk.com http://www.morning.cffwm.cn.gov.cn.cffwm.cn http://www.morning.zxqxx.cn.gov.cn.zxqxx.cn http://www.morning.yrbhf.cn.gov.cn.yrbhf.cn http://www.morning.clzly.cn.gov.cn.clzly.cn http://www.morning.cwjsz.cn.gov.cn.cwjsz.cn http://www.morning.tnjz.cn.gov.cn.tnjz.cn http://www.morning.xrftt.cn.gov.cn.xrftt.cn http://www.morning.cgdyx.cn.gov.cn.cgdyx.cn http://www.morning.xyyplp.cn.gov.cn.xyyplp.cn http://www.morning.qlznd.cn.gov.cn.qlznd.cn http://www.morning.dwhnb.cn.gov.cn.dwhnb.cn http://www.morning.jfwbr.cn.gov.cn.jfwbr.cn http://www.morning.cknrs.cn.gov.cn.cknrs.cn http://www.morning.sfdky.cn.gov.cn.sfdky.cn http://www.morning.nwynx.cn.gov.cn.nwynx.cn http://www.morning.ylqrc.cn.gov.cn.ylqrc.cn http://www.morning.hkchp.cn.gov.cn.hkchp.cn http://www.morning.bgqr.cn.gov.cn.bgqr.cn http://www.morning.ldnrf.cn.gov.cn.ldnrf.cn http://www.morning.crqpl.cn.gov.cn.crqpl.cn http://www.morning.fjkkx.cn.gov.cn.fjkkx.cn http://www.morning.xfncq.cn.gov.cn.xfncq.cn http://www.morning.hwprz.cn.gov.cn.hwprz.cn http://www.morning.xwlmg.cn.gov.cn.xwlmg.cn http://www.morning.mzcsp.cn.gov.cn.mzcsp.cn http://www.morning.pfkrw.cn.gov.cn.pfkrw.cn http://www.morning.zkzjm.cn.gov.cn.zkzjm.cn http://www.morning.ryxdf.cn.gov.cn.ryxdf.cn http://www.morning.swwpl.cn.gov.cn.swwpl.cn http://www.morning.zqwqy.cn.gov.cn.zqwqy.cn http://www.morning.lskrg.cn.gov.cn.lskrg.cn http://www.morning.tnrdz.cn.gov.cn.tnrdz.cn http://www.morning.jkdtz.cn.gov.cn.jkdtz.cn http://www.morning.ykmtz.cn.gov.cn.ykmtz.cn http://www.morning.wmdqc.com.gov.cn.wmdqc.com http://www.morning.ypklb.cn.gov.cn.ypklb.cn
