torch.utils.data

整体架构

平时使用 pytorch 加载数据时大概是这样的:

py 复制代码
import numpy as np
from torch.utils.data import Dataset, DataLoader

class ExampleDataset(Dataset):
	def __init__(self):
		self.data = [1, 2, 3, 4, 5]

	def __getitem__(self, idx):
		return self.data[idx]

	def __len__(self):
		return len(self.data)

def collate_fn(batch):
	return np.array(batch)

dataset = ExampleDataset()  # create the dataset
dataloader = DataLoader(
	dataset=dataset,
	batch_size=2,
	shuffle=True,
	num_workers=4,
	collate_fn=collate_fn
)
for datapoint in dataloader:
	print(datapoint)
  1. 继承 Dataset 类,定义一个迭代器,包含两个魔法方法:__getitem__(self, idx)__len__(self),分别实现如何获取一条数据和如何设定数据长度;
  2. 定义 collate_fn 函数,设定如何组织一个 batch
  3. 实例化 Dataset,并和 collate_fn 一起传入 DataLoader,参数 batch_size 设置批大小、shuffle 设置是否打乱、num_workers 设置并行加载数据的进程数。

然而,背后到底干了什么,我们不清楚,甚至遇到 DataLoader 的如 samplerbatch_samplerworker_init_fn 的其他参数,就会懵逼。那就看一看官方文档,了解一下 torch.utils.data 是如何工作的。

上图是数据加载的整体框架图,官网说 DataLoader 组合datasetsampler,多个 workers 根据 dataset 提供的数据副本sampler 提供的 keys 并行地加载数据,并通过 collate_fn 组成 batch 供用户迭代。需要注意的有:

  1. 每个 worker 持有数据的一个副本,故占用内存 "主线程内存 * num_workers";
  2. 即使用户不提供 sampler 对象 (通常不提供),DataLoader 也会根据 shuffle 参数创建一个默认的 sampler 对象 ;一旦提供了,其前路的 shuffle 参数不能为 True (不提供就好);
  3. 即使用户不提供 batch_sampler 对象 (通常不提供),DataLoader 也会根据 batch_sampler, drop_last 参数创建一个默认的 batch_sampler 对象 ;一旦提供了,其前路的 shuffle, drop_last 不能为 Truebatch_size 必须为 1 1 1,sampler 必须为 None,因为创建 BatchSampler 时已经有了这些参数;

    本质上是把创建 batch_sampler 的活拉出来由用户在 DataLoader 外自定义地做了。

Dataset

分为两种:map-styleiterable-style 。前者的数据可通过 [idx or key] 访问,后者的数据只能通过迭代器 next 一个个访问。所以上面架构中的采样器是对于 map-style 数据集说的iterable-style 的数据集的访问顺序由迭代器决定。

Sampler

torch.utils.data.Sampler 的子类或 Iterable,两个例子:

py 复制代码
class AccedingSequenceLengthSampler(tu_data.Sampler[int]):
	def __init__(self, data: List[str]) -> None:
		super().__init__()
		self.data = data

	def __len__(self) -> int:
		return len(self.data)

	def __iter__(self) -> Iterator[int]:
		"""
		:return: 实现了按数据长短顺序访问数据集
		"""
		sizes = torch.tensor([len(x) for x in self.data])
		yield from torch.argsort(sizes).tolist()


class AccedingSequenceLengthBatchSampler(tu_data.Sampler[List[int]]):
	def __init__(self, data: List[str], batch_size: int) -> None:
		super().__init__()
		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()

Batch

batch_sampler 提供一批下标,取得一批数据后由 collate_fn 将这批数据整合:

py 复制代码
if collate_fn is None:
	if self._auto_collation:
		collate_fn = _utils.collate.default_collate
	else:  # self.batch_sampler is None: (batch_size is None) and (batch_sampler is None)
		collate_fn = _utils.collate.default_convert

分两种情况:

  • automatic batching is disabled :调用 default_convert 函数简单地将 NumPy arrays 转化为 PyTorch Tensor;
  • automatic batching is enabled :调用 default_collate 函数,转化会变得复杂一点:
py 复制代码
from torch.utils import data as tu_data
import collections

# %% Example with a batch of `int`s:
tu_data.default_collate([0, 1, 2, 3])
# tensor([0, 1, 2, 3])

# %% Example with a batch of `str`s:
tu_data.default_collate(['a', 'b', 'c'])
# ['a', 'b', 'c']

# %% Example with `Map` inside the batch:
tu_data.default_collate([
	{'A': 0, 'B': 1},
	{'A': 100, 'B': 100}
])
# {'A': tensor([0, 100]), 'B': tensor([1, 100])}, 同 key 的合并了

# %% Example with `NamedTuple` inside the batch:
Point = collections.namedtuple('Point', ['x', 'y'])
tu_data.default_collate([Point(0, 0), Point(1, 1)])
# Point(x=tensor([0, 1]), y=tensor([0, 1])), 同 name 的合并了, 大概和 dict 一样吧

# %% Example with `Tuple` inside the batch:
tu_data.default_collate([(0, 1), (2, 3)])
# [tensor([0, 2]), tensor([1, 3])], 对 list 内部执行 collate

# %% Example with `List` inside the batch:
tu_data.default_collate([[0, 1], [2, 3]])  # [tensor([0, 2]), tensor([1, 3])], 对 list 内部执行 collate, 并没有变成二维 tensor

Multi-process Data Loading

dataset, collate_fn, and worker_init_fn are passed to each worker ,大概能说明 batch 是在子进程内部 合成的。

有一个需要注意的地方是内存增长问题,当 __get_item__(self, key) 访问数据时,由于 Python 对象的 refcount 机制,数据会不断地复制 ,从而内存爆炸。但这里说解决 number of workers * size of parent process 问题,就不追究了,反正尽量用 numpy 或 pytorch tensor 吧。
iterable-style datasets 的随机性

相关推荐
仗剑_走天涯5 小时前
基于pytorch.nn模块实现线性模型
人工智能·pytorch·python·深度学习
2401_8786247912 小时前
pytorch 自动微分
人工智能·pytorch·python·机器学习
水龙吟啸14 小时前
从零开始搭建深度学习大厦系列-2.卷积神经网络基础(5-9)
人工智能·pytorch·深度学习·cnn·mxnet
慕婉030717 小时前
深度学习中的常见损失函数详解及PyTorch实现
人工智能·pytorch·深度学习
聚客AI18 小时前
搜索引擎vs向量数据库:LangChain混合检索架构实战解析
人工智能·pytorch·语言模型·自然语言处理·数据分析·gpt-3·文心一言
咸鱼鲸20 小时前
【PyTorch】PyTorch中数据准备工作(AI生成)
人工智能·pytorch·python
羊八井20 小时前
使用 Earth2Studio 和 AI 模型进行全球天气预测:太阳辐照
pytorch·python·nvidia
向左转, 向右走ˉ21 小时前
PyTorch随机擦除:提升模型抗遮挡能力
人工智能·pytorch·python·深度学习
HuashuiMu花水木2 天前
PyTorch笔记3----------统计学相关函数
人工智能·pytorch·笔记
AndrewHZ2 天前
【图像处理基石】如何检测到画面中的ppt并对其进行增强?
图像处理·人工智能·pytorch·opencv·目标检测·计算机视觉·图像增强