torch.utils.data

Типы наборов данных

Самый важный аргумент конструктора DataLoader — dataset, который указывает объект набора данных, из которого нужно загружать данные. PyTorch поддерживает два типа наборов данных:

• Наборы данных с отображением,
• Итеративные наборы данных.

Порядок загрузки данных и сэмплер

Для итеративных наборов данных порядок загрузки данных полностью определяется пользовательским итерируемым объектом. Это упрощает реализацию чтения блоками и динамического размера пакета (например, выдавая пакетный образец при каждой итерации).

Остальная часть этого раздела посвящена случаю с наборами данных с отображением. Классы torch.utils.data.Sampler используются для задания последовательности индексов/ключей, применяемых при загрузке данных. Они представляют собой итерируемые объекты, перебирающие индексы наборов данных. Например, в распространённом случае стохастического градиентного спуска (SGD) объект Sampler может случайным образом переставить список индексов и выдавать их по одному за раз либо выдавать несколько индексов для мини-пакетного SGD.

Последовательный или перемешанный сэмплер будет автоматически создан на основе аргумента shuffle объекта DataLoader. Также пользователи могут использовать аргумент sampler, чтобы указать пользовательский объект Sampler, который при каждой итерации выдаёт следующий индекс/ключ для получения данных.

В качестве аргумента batch_sampler можно передать пользовательский объект Sampler, который за раз выдаёт список индексов для пакета. Автоматическое формирование пакетов также можно включить с помощью аргументов batch_size и drop_last. Подробнее см. в следующем разделе.

Загрузка пакетных и непакетных данных

DataLoader поддерживает автоматическое объединение полученных отдельных образцов данных в пакеты с помощью аргументов batch_size, drop_last, batch_sampler и collate_fn (для которого предусмотрена функция по умолчанию).

for indices in batch_sampler:
    yield collate_fn([dataset[i] for i in indices])

dataset_iter = iter(dataset)
for indices in batch_sampler:
    yield collate_fn([next(dataset_iter) for _ in indices])

for index in sampler:
    yield collate_fn(dataset[index])

for data in iter(dataset):
    yield collate_fn(data)

Загрузка данных в одном и нескольких процессах

По умолчанию DataLoader загружает данные в одном процессе.

Внутри процесса Python глобальная блокировка интерпретатора (GIL) не позволяет в полной мере распараллеливать выполнение кода Python между потоками. Чтобы загрузка данных не блокировала вычисления, PyTorch позволяет легко включить загрузку данных в нескольких процессах: для этого достаточно задать аргументу num_workers положительное целое число.

Закрепление памяти

Копирование данных с хоста на GPU происходит гораздо быстрее, если данные находятся в закреплённой (заблокированной страницами) памяти. Подробнее о том, когда и как использовать закреплённые буферы памяти в целом, см. в разделе Использование буферов закреплённой памяти.

При загрузке данных передача pin_memory=True в DataLoader автоматически помещает полученные тензоры данных в закреплённую память и тем самым ускоряет передачу данных на GPU с поддержкой CUDA.

Логика закрепления памяти по умолчанию распознаёт только тензоры, а также отображения и итерируемые объекты, содержащие тензоры. По умолчанию, если логика закрепления получает пакет пользовательского типа (это происходит, если у вас есть collate_fn, возвращающий пакет пользовательского типа) или если каждый элемент пакета имеет пользовательский тип, она не распознаёт их и возвращает пакет (или эти элементы) без закрепления памяти. Чтобы включить закрепление памяти для пакета или типов данных пользователя, определите метод pin_memory() в соответствующих пользовательских типах.

См. пример ниже.

Пример:

class SimpleCustomBatch:
    def __init__(self, data):
        transposed_data = list(zip(*data))
        self.inp = torch.stack(transposed_data[0], 0)
        self.tgt = torch.stack(transposed_data[1], 0)

    # custom memory pinning method on custom type
    def pin_memory(self):
        self.inp = self.inp.pin_memory()
        self.tgt = self.tgt.pin_memory()
        return self

def collate_wrapper(batch):
    return SimpleCustomBatch(batch)

inps = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
tgts = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
dataset = TensorDataset(inps, tgts)

loader = DataLoader(dataset, batch_size=2, collate_fn=collate_wrapper,
                    pin_memory=True)

for batch_ndx, sample in enumerate(loader):
    print(sample.inp.is_pinned())
    print(sample.tgt.is_pinned())

class torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=None, sampler=None, batch_sampler=None, num_workers=0, collate_fn=None, pin_memory=False, drop_last=False, timeout=0, worker_init_fn=None, multiprocessing_context=None, generator=None, *, prefetch_factor=None, persistent_workers=False, pin_memory_device='', in_order=True) [исходный код]

Загрузчик данных объединяет набор данных и сэмплер и предоставляет итерируемый объект для заданного набора данных.

DataLoader поддерживает наборы данных как с отображением, так и с итерируемым стилем, загрузку в одном или нескольких процессах, настройку порядка загрузки, а также необязательное автоматическое формирование пакетов (сборку) и закрепление памяти.

Подробнее см. на странице документации torch.utils.data.

Параметры:

• dataset (Dataset) – набор данных, из которого загружаются данные.
• batch_size (int, необязательно) – количество образцов в загружаемом пакете (по умолчанию: 1).
• shuffle (bool, необязательно) – установите значение True, чтобы перемешивать данные в каждой эпохе (по умолчанию: False).
• sampler (Sampler или Iterable, необязательно) – определяет стратегию выборки образцов из набора данных. Может быть любым Iterable с реализованным __len__. Если задано, параметр shuffle задавать нельзя.
• batch_sampler (Sampler или Iterable, необязательно) – аналог sampler, но за один раз возвращает пакет индексов. Не может использоваться одновременно с batch_size, shuffle, sampler и drop_last.
• num_workers (int, необязательно) – количество подпроцессов для загрузки данных. 0 означает, что данные будут загружаться в основном процессе. (по умолчанию: 0)
• collate_fn (Callable, необязательно) – объединяет список образцов, формируя мини-пакет из тензора или тензоров. Используется при пакетной загрузке из набора данных с отображением.
• pin_memory (bool, необязательно) – если задано значение True, загрузчик данных перед возвратом скопирует тензоры в закреплённую память устройства/CUDA. Если элементы данных имеют пользовательский тип или collate_fn возвращает пакет пользовательского типа, см. пример ниже.
• drop_last (bool, необязательно) – установите значение True, чтобы отбросить последний неполный пакет, если размер набора данных не делится на размер пакета. Если задано False, а размер набора данных не делится на размер пакета, последний пакет будет меньше. (по умолчанию: False)
• timeout (числовой тип, необязательно) – если значение положительное, оно задаёт время ожидания получения пакета от рабочих процессов. Значение всегда должно быть неотрицательным. (по умолчанию: 0)
• worker_init_fn (Callable, необязательно) – если значение не равно None, эта функция будет вызвана для каждого рабочего подпроцесса с идентификатором рабочего процесса (целым числом из [0, num_workers - 1]) в качестве аргумента после инициализации начального значения и до загрузки данных. (по умолчанию: None)
• multiprocessing_context (str или multiprocessing.context.BaseContext, необязательно) – если задано значение None, будет использоваться контекст многопроцессной обработки # noqa: D401 по умолчанию для вашей операционной системы. (по умолчанию: None)
• generator (torch.Generator, необязательно) – если значение не равно None, этот генератор случайных чисел будет использоваться RandomSampler для генерации случайных индексов, а многопроцессной обработкой — для генерации base_seed рабочих процессов. (по умолчанию: None)
• prefetch_factor (int, необязательно, только именованный аргумент) – количество пакетов, предварительно загружаемых каждым рабочим процессом. 2 означает, что всего для всех рабочих процессов будет предварительно загружено 2 * num_workers пакетов. (значение по умолчанию зависит от заданного значения num_workers. Если num_workers=0, значение по умолчанию — None. В противном случае, если значение num_workers > 0, по умолчанию используется 2).
• persistent_workers (bool, необязательно) – если задано значение True, загрузчик данных не будет завершать рабочие процессы после однократного использования набора данных. Это позволяет сохранять экземпляры Dataset рабочих процессов активными. (по умолчанию: False)
• pin_memory_device (str, необязательно) – устарел; если задано значение pin_memory=True, в качестве устройства будет использоваться текущий ускоритель.
• in_order (bool, необязательно) – если задано значение False, загрузчик данных не будет гарантировать возврат пакетов в порядке «первым пришёл — первым обслужен». Применяется только при значении num_workers > 0. (по умолчанию: True)

Предупреждение

Если используется метод запуска spawn, worker_init_fn не может быть объектом, который нельзя сериализовать с помощью pickle, например лямбда-функцией. Подробнее о многопроцессной обработке в PyTorch см. в разделе Рекомендации по многопроцессной обработке.

Предупреждение

Эвристика len(dataloader) основана на длине используемого сэмплера. Если dataset является IterableDataset, вместо этого возвращается оценка на основе len(dataset) / batch_size с соответствующим округлением в зависимости от drop_last, независимо от конфигурации многопроцессной загрузки. Это наилучшая оценка, которую может дать PyTorch, поскольку PyTorch доверяет пользовательскому коду dataset и предполагает, что он правильно обрабатывает многопроцессную загрузку, исключая дублирование данных.

Однако, если при распределении по сегментам у нескольких рабочих процессов остаются неполные последние пакеты, эта оценка всё ещё может быть неточной, поскольку (1) в остальном полный пакет может быть разделён на несколько пакетов и (2) при заданном значении drop_last может быть отброшено больше образцов, чем помещается в один пакет. К сожалению, PyTorch не может в общем случае обнаружить такие ситуации.

Подробнее об этих двух типах наборов данных и о взаимодействии IterableDataset с многопроцессной загрузкой данных см. в разделе Типы наборов данных.

Предупреждение

Ответы на вопросы о начальных значениях генератора случайных чисел см. в разделах Воспроизводимость, Рабочие процессы загрузчика данных возвращают одинаковые случайные числа и Случайность при многопроцессной загрузке данных.

Предупреждение

Установка in_order в False может снизить воспроизводимость и привести к смещённому распределению данных, передаваемых тренеру, в случаях несбалансированных данных.

class torch.utils.data.Dataset [исходный код]

Абстрактный класс, представляющий Dataset.

Все наборы данных, представляющие отображение ключей на образцы данных, должны быть его подклассами. Все подклассы должны переопределять __getitem__(), поддерживая получение образца данных по заданному ключу. Подклассы также могут переопределить __len__(), который, как ожидается, возвращает размер набора данных; он используется многими реализациями Sampler и параметрами по умолчанию DataLoader. Для ускорения пакетной загрузки образцов подклассы также могут реализовать __getitems__(). Этот метод принимает список индексов образцов пакета и возвращает список образцов.

Примечание

DataLoader по умолчанию создаёт сэмплер индексов, выдающий целочисленные индексы. Чтобы он работал с набором данных с отображением и нецелочисленными индексами/ключами, необходимо указать пользовательский сэмплер.

class torch.utils.data.IterableDataset [исходный код]

Итерируемый набор данных.

Все наборы данных, представляющие итерируемую последовательность образцов данных, должны быть его подклассами. Такая форма наборов данных особенно полезна, когда данные поступают из потока.

Все подклассы должны переопределять __iter__(), который возвращает итератор образцов этого набора данных.

При использовании подкласса с DataLoader каждый элемент набора данных будет выдаваться итератором DataLoader. При значении num_workers > 0 каждый рабочий процесс получает отдельную копию объекта набора данных, поэтому часто требуется настраивать каждую копию независимо, чтобы рабочие процессы не возвращали повторяющиеся данные. get_worker_info(), вызванный в рабочем процессе, возвращает информацию о нём. Его можно использовать как в методе __iter__() набора данных, так и в параметре worker_init_fn объекта DataLoader, чтобы изменить поведение каждой копии.

Пример 1: распределение нагрузки между всеми рабочими процессами в __iter__():

>>> class MyIterableDataset(torch.utils.data.IterableDataset):
...     def __init__(self, start, end):
...         super(MyIterableDataset).__init__()
...         assert end > start, "this example only works with end >= start"
...         self.start = start
...         self.end = end
...
...     def __iter__(self):
...         worker_info = torch.utils.data.get_worker_info()
...         if worker_info is None:  # single-process data loading, return the full iterator
...             iter_start = self.start
...             iter_end = self.end
...         else:  # in a worker process
...             # split workload
...             per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers)))
...             worker_id = worker_info.id
...             iter_start = self.start + worker_id * per_worker
...             iter_end = min(iter_start + per_worker, self.end)
...         return iter(range(iter_start, iter_end))
...
>>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].
>>> ds = MyIterableDataset(start=3, end=7)

>>> # Single-process loading
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=0)))
[tensor([3]), tensor([4]), tensor([5]), tensor([6])]

>>> # Multi-process loading with two worker processes
>>> # Worker 0 fetched [3, 4].  Worker 1 fetched [5, 6].
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2)))
[tensor([3]), tensor([5]), tensor([4]), tensor([6])]

>>> # With even more workers
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=12)))
[tensor([3]), tensor([5]), tensor([4]), tensor([6])]

Пример 2: распределение нагрузки между всеми рабочими процессами с помощью worker_init_fn:

>>> class MyIterableDataset(torch.utils.data.IterableDataset):
...     def __init__(self, start, end):
...         super(MyIterableDataset).__init__()
...         assert end > start, "this example only works with end >= start"
...         self.start = start
...         self.end = end
...
...     def __iter__(self):
...         return iter(range(self.start, self.end))
...
>>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].
>>> ds = MyIterableDataset(start=3, end=7)

>>> # Single-process loading
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=0)))
[3, 4, 5, 6]
>>>
>>> # Directly doing multi-process loading yields duplicate data
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2)))
[3, 3, 4, 4, 5, 5, 6, 6]

>>> # Define a `worker_init_fn` that configures each dataset copy differently
>>> def worker_init_fn(worker_id):
...     worker_info = torch.utils.data.get_worker_info()
...     dataset = worker_info.dataset  # the dataset copy in this worker process
...     overall_start = dataset.start
...     overall_end = dataset.end
...     # configure the dataset to only process the split workload
...     per_worker = int(math.ceil((overall_end - overall_start) / float(worker_info.num_workers)))
...     worker_id = worker_info.id
...     dataset.start = overall_start + worker_id * per_worker
...     dataset.end = min(dataset.start + per_worker, overall_end)
...

>>> # Multi-process loading with the custom `worker_init_fn`
>>> # Worker 0 fetched [3, 4].  Worker 1 fetched [5, 6].
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2, worker_init_fn=worker_init_fn)))
[3, 5, 4, 6]

>>> # With even more workers
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=12, worker_init_fn=worker_init_fn)))
[3, 4, 5, 6]

class torch.utils.data.TensorDataset(*tensors) [исходный код]

Набор данных, оборачивающий тензоры.

Каждый образец извлекается путём индексирования тензоров по первому измерению.

Параметры:

*tensors (Tensor) – тензоры с одинаковым размером первого измерения.

class torch.utils.data.StackDataset(*args, **kwargs) [исходный код]

Набор данных, составленный из нескольких наборов данных.

Этот класс полезен для объединения разных частей сложных входных данных, представленных наборами данных.

Пример

>>> images = ImageDataset()
>>> texts = TextDataset()
>>> tuple_stack = StackDataset(images, texts)
>>> tuple_stack[0] == (images[0], texts[0])
>>> dict_stack = StackDataset(image=images, text=texts)
>>> dict_stack[0] == {"image": images[0], "text": texts[0]}

Параметры:

• *args (Dataset) – наборы данных для объединения; результат возвращается в виде кортежа.
• **kwargs (Dataset) – наборы данных для объединения; результат возвращается в виде словаря.

class torch.utils.data.ConcatDataset(datasets) [исходный код]

Набор данных, представляющий собой конкатенацию нескольких наборов данных.

Этот класс полезен для объединения нескольких существующих наборов данных.

Параметры:

datasets (последовательность) – список наборов данных для конкатенации

class torch.utils.data.ChainDataset(datasets) [исходный код]

Набор данных для объединения нескольких IterableDataset.

Этот класс полезен для объединения потоков из разных существующих наборов данных. Операция объединения выполняется на лету, поэтому этот класс обеспечивает эффективную конкатенацию крупномасштабных наборов данных.

Параметры:

datasets (итерируемый объект из IterableDataset) – наборы данных для объединения

class torch.utils.data.Subset(dataset, indices) [исходный код]

Подмножество набора данных с заданными индексами.

Примечание

При создании подкласса Subset и переопределении __getitem__ необходимо также переопределить __getitems__, чтобы обеспечить корректную работу DataLoader с пользовательской логикой. Если переопределить только __getitem__, при использовании DataLoader будет вызвано NotImplementedError.

Простая реализация __getitems__ может делегировать вызов __getitem__:

def __getitems__(self, indices):
    return [self.__getitem__(idx) for idx in indices]

Для повышения производительности рассмотрите возможность реализации логики обработки пакетов в __getitems__ вместо многократного вызова __getitem__.

Параметры:

• dataset (Dataset) – весь набор данных
• indices (последовательность) – индексы элементов полного набора, выбранные для подмножества

torch.utils.data._utils.collate.collate(batch, *, collate_fn_map=None) [исходный код]

Универсальная функция сборки, обрабатывающая тип коллекции элементов в каждом пакете.

Функция также предоставляет реестр функций для обработки определённых типов элементов. default_collate_fn_map содержит функции сборки по умолчанию для тензоров, массивов NumPy, чисел и строк.

Параметры:

• batch – один пакет для сборки
• collate_fn_map (dict[type | tuple[type, ...], Callable] | None) – необязательный словарь, сопоставляющий типы элементов соответствующим функциям сборки. Если тип элемента отсутствует в этом словаре, функция перебирает его ключи в порядке добавления и вызывает соответствующую функцию сборки, если тип элемента является подклассом ключа.

Примеры

>>> def collate_tensor_fn(batch, *, collate_fn_map):
...     # Extend this function to handle batch of tensors
...     return torch.stack(batch, 0)
>>> def custom_collate(batch):
...     collate_map = {torch.Tensor: collate_tensor_fn}
...     return collate(batch, collate_fn_map=collate_map)
>>> # Extend `default_collate` by in-place modifying `default_collate_fn_map`
>>> default_collate_fn_map.update({torch.Tensor: collate_tensor_fn})

Примечание

Каждая функция сборки требует позиционный аргумент для пакета и именованный аргумент для словаря функций сборки в качестве collate_fn_map.

torch.utils.data.default_collate(batch) [исходный код]

Принимает пакет данных и помещает его элементы в тензор с дополнительным внешним измерением — размером пакета.

Точный тип результата может быть torch.Tensor, Sequence из torch.Tensor, коллекцией torch.Tensor или оставаться без изменений — в зависимости от типа входных данных. Эта функция используется для сборки по умолчанию, если в DataLoader заданы batch_size или batch_sampler.

Ниже приведено соответствие общего типа входных данных (определяемого типом элемента пакета) и типа выходных данных:

• torch.Tensor -> torch.Tensor (с добавленным внешним измерением — размером пакета)
• Массивы NumPy -> torch.Tensor
• float -> torch.Tensor
• int -> torch.Tensor
• str -> str (без изменений)
• bytes -> bytes (без изменений)
• Mapping[K, V_i] -> Mapping[K, default_collate([V_1, V_2, …])]
• NamedTuple[V1_i, V2_i, …] -> NamedTuple[default_collate([V1_1, V1_2, …]), default_collate([V2_1, V2_2, …]), …]
• Sequence[V1_i, V2_i, …] -> Sequence[default_collate([V1_1, V1_2, …]), default_collate([V2_1, V2_2, …]), …]

Параметры:

batch – один пакет для сборки

Примеры

>>> # Example with a batch of `int`s:
>>> default_collate([0, 1, 2, 3])
tensor([0, 1, 2, 3])
>>> # Example with a batch of `str`s:
>>> default_collate(["a", "b", "c"])
['a', 'b', 'c']
>>> # Example with `Map` inside the batch:
>>> default_collate([{"A": 0, "B": 1}, {"A": 100, "B": 100}])
{'A': tensor([  0, 100]), 'B': tensor([  1, 100])}
>>> # Example with `NamedTuple` inside the batch:
>>> Point = namedtuple("Point", ["x", "y"])
>>> default_collate([Point(0, 0), Point(1, 1)])
Point(x=tensor([0, 1]), y=tensor([0, 1]))
>>> # Example with `Tuple` inside the batch:
>>> default_collate([(0, 1), (2, 3)])
[tensor([0, 2]), tensor([1, 3])]
>>> # Example with `List` inside the batch:
>>> default_collate([[0, 1], [2, 3]])
[tensor([0, 2]), tensor([1, 3])]
>>> # Two options to extend `default_collate` to handle specific type
>>> # Option 1: Write custom collate function and invoke `default_collate`
>>> def custom_collate(batch):
...     elem = batch[0]
...     if isinstance(elem, CustomType):  # Some custom condition
...         return ...
...     else:  # Fall back to `default_collate`
...         return default_collate(batch)
>>> # Option 2: In-place modify `default_collate_fn_map`
>>> def collate_customtype_fn(batch, *, collate_fn_map=None):
...     return ...
>>> default_collate_fn_map.update(CustomType, collate_customtype_fn)
>>> default_collate(batch)  # Handle `CustomType` automatically

torch.utils.data.default_convert(data) [исходный код]

Преобразует каждый элемент-массив NumPy в torch.Tensor.

Если входные данные — Sequence, Collection или Mapping, функция пытается преобразовать каждый элемент внутри в torch.Tensor. Если входные данные не являются массивом NumPy, они остаются без изменений. Эта функция используется для сборки по умолчанию, если в DataLoader НЕ заданы ни batch_sampler, ни batch_size.

Общее соответствие типов входных и выходных данных аналогично default_collate(). Подробнее см. описание этой функции.

Параметры:

data – одна точка данных для преобразования

Примеры

>>> # Example with `int`
>>> default_convert(0)
0
>>> # Example with NumPy array
>>> default_convert(np.array([0, 1]))
tensor([0, 1])
>>> # Example with NamedTuple
>>> Point = namedtuple("Point", ["x", "y"])
>>> default_convert(Point(0, 0))
Point(x=0, y=0)
>>> default_convert(Point(np.array(0), np.array(0)))
Point(x=tensor(0), y=tensor(0))
>>> # Example with List
>>> default_convert([np.array([0, 1]), np.array([2, 3])])
[tensor([0, 1]), tensor([2, 3])]

torch.utils.data.get_worker_info() [исходный код]

Возвращает сведения о текущем рабочем процессе итератора DataLoader.

При вызове в рабочем процессе функция возвращает объект, который гарантированно имеет следующие атрибуты:

• id: идентификатор текущего рабочего процесса.
• num_workers: общее число рабочих процессов.
• seed: начальное значение генератора случайных чисел для текущего рабочего процесса. Оно определяется генератором случайных чисел основного процесса и идентификатором рабочего процесса. Подробнее см. документацию DataLoader.
• dataset: копия объекта набора данных в этом процессе. Обратите внимание, что в другом процессе это будет другой объект, отличный от объекта основного процесса.

При вызове в основном процессе функция возвращает None.

Примечание

При использовании в worker_init_fn, переданном в DataLoader, этот метод может быть полезен для разной настройки каждого рабочего процесса: например, с помощью worker_id можно настроить объект dataset так, чтобы он читал только определённую часть набора данных, разбитого на сегменты, или использовать seed для инициализации генераторов случайных чисел других библиотек, используемых в коде набора данных.

Тип возвращаемого значения:

WorkerInfo | None

torch.utils.data.dataset.random_split(dataset, lengths, generator=<torch._C.Generator object>) [исходный код]

Случайным образом разделяет набор данных на непересекающиеся новые наборы данных заданной длины.

Если задан список долей, сумма которых равна 1, длины будут автоматически вычислены как floor(frac * len(dataset)) для каждой указанной доли.

После вычисления длин, если останутся элементы, по одному элементу будет распределяться по длинам циклически, пока остатки не закончатся.

При необходимости зафиксируйте генератор для воспроизводимых результатов, например:

Пример

>>> generator1 = torch.Generator().manual_seed(42)
>>> generator2 = torch.Generator().manual_seed(42)
>>> random_split(range(10), [3, 7], generator=generator1)
>>> random_split(range(30), [0.3, 0.3, 0.4], generator=generator2)

Параметры:

• dataset (Dataset) – набор данных для разделения
• lengths (последовательность) – длины или доли создаваемых частей
• generator (Generator) – генератор, используемый для случайной перестановки.

Тип возвращаемого значения:

list[Subset[_T]]

class torch.utils.data.Sampler [исходный код]

Базовый класс для всех сэмплеров.

Каждый подкласс Sampler должен предоставлять метод __iter__(), позволяющий перебирать индексы или списки индексов (пакеты) элементов набора данных, и может предоставлять метод __len__(), который возвращает длину возвращаемых итераторов.

Пример

>>> class AccedingSequenceLengthSampler(Sampler[int]):
>>>     def __init__(self, data: List[str]) -> None:
>>>         self.data = data
>>>
>>>     def __len__(self) -> int:
>>>         return len(self.data)
>>>
>>>     def __iter__(self) -> Iterator[int]:
>>>         sizes = torch.tensor([len(x) for x in self.data])
>>>         yield from torch.argsort(sizes).tolist()
>>>
>>> class AccedingSequenceLengthBatchSampler(Sampler[List[int]]):
>>>     def __init__(self, data: List[str], batch_size: int) -> None:
>>>         self.data = data
>>>         self.batch_size = batch_size
>>>
>>>     def __len__(self) -> int:
>>>         return (len(self.data) + self.batch_size - 1) // self.batch_size
>>>
>>>     def __iter__(self) -> Iterator[List[int]]:
>>>         sizes = torch.tensor([len(x) for x in self.data])
>>>         for batch in torch.chunk(torch.argsort(sizes), len(self)):
>>>             yield batch.tolist()

Примечание

Метод __len__() не является строго обязательным для DataLoader, но ожидается при любых вычислениях, связанных с длиной DataLoader.

class torch.utils.data.SequentialSampler(data_source) [исходный код]

Последовательно выбирает элементы, всегда в одном и том же порядке.

Параметры:

data_source (Sized) – источник данных для выборки. Должен реализовывать __len__.

class torch.utils.data.RandomSampler(data_source, replacement=False, num_samples=None, generator=None) [исходный код]

Случайным образом выбирает элементы. Если выборка выполняется без возвращения, элементы выбираются из перемешанного набора данных.

При выборке с возвращением пользователь может указать num_samples выбираемых элементов.

Параметры:

• data_source (Sized) – источник данных для выборки. Должен реализовывать __len__.
• replacement (bool) – если True, выборка выполняется по запросу с возвращением; по умолчанию ``False``
• num_samples (int) – число выбираемых элементов; по умолчанию `len(dataset)`.
• generator (Generator) – генератор, используемый при выборке.

class torch.utils.data.SubsetRandomSampler(indices, generator=None) [исходный код]

Случайным образом выбирает элементы из заданного списка индексов без возвращения.

Параметры:

• indices (sequence) – последовательность индексов
• generator (Generator) – генератор, используемый при выборке.

class torch.utils.data.WeightedRandomSampler(weights, num_samples, replacement=True, generator=None) [исходный код]

Выбирает элементы из [0,..,len(weights)-1] с заданными вероятностями (весами).

Параметры:

• weights (sequence) – последовательность весов, сумма которых не обязательно равна единице
• num_samples (int) – число выбираемых элементов
• replacement (bool) – если True, выборка выполняется с возвращением. В противном случае она выполняется без возвращения, то есть после выбора индекса элемента для строки его нельзя выбрать для этой строки повторно.
• generator (Generator) – генератор, используемый при выборке.

Пример

>>> list(
...     WeightedRandomSampler(
...         [0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=True
...     )
... )
[4, 4, 1, 4, 5]
>>> list(
...     WeightedRandomSampler(
...         [0.9, 0.4, 0.05, 0.2, 0.3, 0.1], 5, replacement=False
...     )
... )
[0, 1, 4, 3, 2]

class torch.utils.data.BatchSampler(sampler, batch_size, drop_last) [исходный код]

Оборачивает другой сэмплер, чтобы формировать мини-пакеты индексов.

Параметры:

• sampler (Sampler or Iterable) – базовый сэмплер. Может быть любым итерируемым объектом
• batch_size (int) – размер мини-пакета.
• drop_last (bool) – если True, сэмплер отбросит последний пакет, если его размер меньше batch_size

Пример

>>> list(
...     BatchSampler(
...         SequentialSampler(range(10)), batch_size=3, drop_last=False
...     )
... )
[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]
>>> list(
...     BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=True)
... )
[[0, 1, 2], [3, 4, 5], [6, 7, 8]]

class torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=None, rank=None, shuffle=True, seed=0, drop_last=False) [исходный код]

Сэмплер, ограничивающий загрузку данных подмножеством набора данных.

Он особенно полезен в сочетании с torch.nn.parallel.DistributedDataParallel. В таком случае каждый процесс может передать экземпляр DistributedSampler в качестве сэмплера для DataLoader и загрузить часть исходного набора данных, предназначенную только для него.

Примечание

Предполагается, что размер набора данных постоянен и что любой его экземпляр всегда возвращает одни и те же элементы в одном и том же порядке.

Параметры:

• dataset (Dataset) – набор данных, используемый для выборки.
• num_replicas (int, optional) – число процессов, участвующих в распределённом обучении. По умолчанию world_size получается из текущей распределённой группы.
• rank (int, optional) – ранг текущего процесса в num_replicas. По умолчанию rank получается из текущей распределённой группы.
• shuffle (bool, optional) – если True (по умолчанию), сэмплер будет перемешивать индексы.
• seed (int, optional) – начальное значение генератора случайных чисел, используемое для перемешивания сэмплера, если shuffle=True. Это число должно быть одинаковым для всех процессов в распределённой группе. По умолчанию: 0.
• drop_last (bool, optional) – если True, сэмплер отбросит остаток данных, чтобы их количество равномерно делилось на число реплик. Если False, сэмплер добавит дополнительные индексы, чтобы количество данных равномерно делилось на число реплик. По умолчанию: False.

Предупреждение

В распределённом режиме для корректной работы перемешивания в нескольких эпохах необходимо вызывать метод set_epoch() в начале каждой эпохи до создания итератора DataLoader. В противном случае всегда будет использоваться один и тот же порядок.

Пример:

>>> sampler = DistributedSampler(dataset) if is_distributed else None
>>> loader = DataLoader(dataset, shuffle=(sampler is None),
...                     sampler=sampler)
>>> for epoch in range(start_epoch, n_epochs):
...     if is_distributed:
...         sampler.set_epoch(epoch)
...     train(loader)

torch.utils.data.graph.traverse(datapipe, only_datapipe=None) [исходный код]

Обходит DataPipes и их атрибуты, чтобы извлечь граф DataPipe.

[Устарела] Если only_dataPipe задано как True, рассматриваются только атрибуты каждой DataPipe, которые сами являются DataPipe или объектами-коллекциями Python, такими как list, tuple, set и dict.

Примечание

Эта функция устарела. Вместо неё используйте traverse_dps.

Параметры:

• datapipe (IterDataPipe | MapDataPipe) – конечная DataPipe графа
• only_datapipe (bool | None) – если False (по умолчанию), обходятся все атрибуты каждой DataPipe. Этот аргумент устарел и будет удалён после следующего выпуска.

Возвращает:

Граф в виде вложенного словаря, где ключи — это идентификаторы экземпляров DataPipe, а значения — кортежи из экземпляра DataPipe и подграфа

Тип возвращаемого значения:

dict[int, tuple[IterDataPipe | MapDataPipe, DataPipeGraph]]

torch.utils.data.graph.traverse_dps(datapipe) [исходный код]

Обходит DataPipes и их атрибуты, чтобы извлечь граф DataPipe.

Рассматриваются только атрибуты каждой DataPipe, которые сами являются DataPipe или объектами-коллекциями Python, такими как list, tuple, set и dict.

Параметры:

datapipe (IterDataPipe | MapDataPipe) – конечная DataPipe графа

Возвращает:

Граф в виде вложенного словаря, где ключи — это идентификаторы экземпляров DataPipe, а значения — кортежи из экземпляра DataPipe и подграфа

Тип возвращаемого значения:

dict[int, tuple[IterDataPipe | MapDataPipe, DataPipeGraph]]

torch.utils.data.datapipes.utils.decoder.basichandlers(extension, data) [исходный код]

Преобразует исходные данные (поток байтов) в объекты Python.

Определяет расширение и загружает данные в объект Python, поддерживающий соответствующее расширение.

Параметры:

• extension (str) – расширение файла
• data (byte stream) – данные для загрузки в объект Python.

Возвращает:

Данные, загруженные в соответствующий объект Python,

поддерживающий это расширение.

Тип возвращаемого значения:

object

Пример

>>> import pickle
>>> data = pickle.dumps("some data")
>>> new_data = basichandlers("pickle", data)
>>> new_data
some data

Преобразование данных для расширений:

• txt, text, transcript: данные в формате str, декодированные из utf-8
• cls, cls2, class, count, index, inx, id: int
• json, jsn: данные, загруженные из json
• pickle, pyd: данные, загруженные из pickle
• pt: данные, загруженные с помощью torch

torch.utils.data.datapipes.utils.decoder.handle_extension(extensions, f) [исходный код]

Возвращает функцию-обработчик декодера для списка расширений.

Расширения могут быть разделённым пробелами списком. Они могут содержать точки; в этом случае ключ, передаваемый в f, должен содержать соответствующее число компонентов расширения. Сравнение не учитывает регистр. Примеры: handle_extension(“jpg jpeg”, my_decode_jpg) # вызывается для любого файла.jpg handle_extension(“seg.jpg”, special_case_jpg) # вызывается только для файла.seg.jpg

torch.utils.data.datapipes.utils.common.validate_input_col(fn, input_col) [исходный код]

Проверяет, что функция, используемая в вызываемой DataPipe, работает с входным столбцом.

Она просто проверяет, совпадает ли число позиционных аргументов с размером входного столбца. Функция не должна содержать обязательных аргументов, передаваемых только по ключевому слову.

Примеры

>>> def f(a, b, *, c=1):
>>>     return a + b + c
>>> def f_def(a, b=1, *, c=1):
>>>     return a + b + c
>>> assert validate_input_col(f, [1, 2])
>>> assert validate_input_col(f_def, 1)
>>> assert validate_input_col(f_def, [1, 2])

Примечания

Если функция содержит переменное число позиционных аргументов (inspect.VAR_POSITIONAL), например f(a, *args), валидатор примет любой размер входного столбца, не меньший числа позиционных аргументов (в данном случае — 1).

Параметры:

• fn (Callable) – функция для проверки.
• input_col (int | tuple | list | None) – входной столбец для проверки.

Вызывает исключение:

ValueError – если функция несовместима с входным столбцом.