torch.utils.data

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

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

• стилизированные наборы данных,
• итерируемые наборы данных.

Порядок загрузки данных и выборка

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

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

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

Настраиваемый Sampler, который выдаёт список индексов пакета за раз, может быть передан в качестве аргумента batch_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 автоматически поместит полученные тензоры данных в фиксированную память, что позволит ускорить передачу данных на CUDA-совместимые GPU.

Логика по умолчанию по фиксированию памяти распознает только тензоры и отображения и итерируемые объекты, содержащие тензоры. По умолчанию, если логика фиксирования видит пакет пользовательского типа (что произойдет, если у вас есть 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=2, persistent_workers=False, pin_memory_device='') [source]

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

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).
• и т.д. (остальные параметры)

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

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

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

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

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

См. Типы наборов данных для получения дополнительной информации об этих двух типах наборов данных и о том, как IterableDataset взаимодействует с Многопроцессорная загрузка данных.

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

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

class torch.utils.data.Dataset(*args, **kwds) [source]

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

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

Примечание

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

class torch.utils.data.IterableDataset(*args, **kwds) [source]

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

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

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

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

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

>>> class MyIterableDataset(torch.utils.data.IterableDataset):
...     def __init__(self, start, end):
...         super(MyIterableDataset).__init__()
...         assert end > start, "this example code 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])]

>>> # Mult-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 code 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)
...

>>> # Mult-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) [source]

Набор данных, содержащий тензоры.

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

Parameters:

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

class torch.utils.data.ConcatDataset(datasets) [source]

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

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

Parameters:

datasets (sequence) – Список наборов данных, которые необходимо конкатенировать

class torch.utils.data.ChainDataset(datasets) [source]

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

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

Parameters:

datasets (iterable of IterableDataset) – наборы данных, которые необходимо объединить в цепочку

class torch.utils.data.Subset(dataset, indices) [source]

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

Parameters:

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

torch.utils.data._utils.collate.collate(batch, *, collate_fn_map=None) [source]

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

Parameters:

• batch – один набор данных для объединения
• collate_fn_map (Optional[Dict[Union[Type, Tuple[Type, ...]], Callable]]) – Необязательный словарь, сопоставляющий тип элемента с соответствующей функцией объединения. Если тип элемента отсутствует в этом словаре, эта функция пройдёт по каждому ключу словаря в порядке вставки, чтобы вызвать соответствующую функцию объединения, если тип элемента является подклассом ключа.

Примеры

>>> # Extend this function to handle batch of tensors
>>> def collate_tensor_fn(batch, *, collate_fn_map):
...     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) [source]

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

Вот общее сопоставление типа входных данных (на основе типа элемента в наборе) и типа выходных данных:

• 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, …]), …]

Parameters:

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(CustoType, collate_customtype_fn)
>>> default_collate(batch)  # Handle `CustomType` automatically

torch.utils.data.default_convert(data) [source]

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

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

Parameters:

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() [source]

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

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

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

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

Примечание

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

torch.utils.data.random_split(dataset, lengths, generator=<torch._C.Generator object>) [source]

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

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

После вычисления длин, если есть остатки, 1 единица будет распределена по длине в порядке циклического обхода, пока не останется остатков.

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

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

Параметры:

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

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

Список[Subset[T]]

class torch.utils.data.Sampler(data_source) [source]

Базовый класс для всех Samplers.

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

Примечание

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

class torch.utils.data.SequentialSampler(data_source) [source]

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

Параметры:

data_source (Dataset) – набор данных для выборки

class torch.utils.data.RandomSampler(data_source, replacement=False, num_samples=None, generator=None) [source]

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

Параметры:

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

class torch.utils.data.SubsetRandomSampler(indices, generator=None) [source]

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

Параметры:

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

class torch.utils.data.WeightedRandomSampler(weights, num_samples, replacement=True, generator=None) [source]

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

Параметры:

• weights (последовательность) – последовательность весов, не обязательно суммирующихся до единицы
• 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) [source]

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

Параметры:

• sampler (Sampler или Iterable) – Базовый sampler. Может быть любым итерируемым объектом
• batch_size (int) – Размер мини-пакета.
• drop_last (bool) – Если True, sampler пропустит последнюю партию, если ее размер будет меньше 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) [source]

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

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

Примечание

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

Параметры:

• dataset – Набор данных, используемый для выборки.
• num_replicas (int, optional) – Количество процессов, участвующих в распределённом обучении. По умолчанию извлекается из текущей распределённой группы.
• rank (int, optional) – Ранг текущего процесса в num_replicas. По умолчанию извлекается из текущей распределённой группы.
• 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)