Своя сборка пакета в PyTorch

При выдаче пакетов DataLoader по умолчанию складывает примеры в список. Иногда нужен другой вид: один тензор, словарь полей или заранее выбранная форма.

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

import torch
from torch.utils.data import DataLoader, TensorDataset

def stack_batch(batch):
	nums = torch.tensor([item[0] for item in batch])
	return nums

data = TensorDataset(torch.tensor([1., 2., 3., 4.]))
loader = DataLoader(
	data, batch_size=2, collate_fn=stack_batch,
)
for batch in loader:
	print(batch)

Создайте набор из чисел [1., 2., 3., 4., 5.] и обход с размером пакета 3. Пусть сборка пакета возвращает один тензор из всех чисел пакета. Выведите форму первого пакета.

Для набора [10., 20., 30., 40.] задайте размер пакета 2 и функцию, которая склеивает примеры в тензор. Выведите содержимое второго пакета.

Создайте набор из шести целых от 0 до 5 и обход с пакетом по 4 примера. Пусть из примеров пакета получается один тензор. Выведите число пакетов за полный проход.