共计 9458 个字符，预计需要花费 24 分钟才能阅读完成。
深度神经网络须要很长时间来训练。训练速度受模型的复杂性、批大小、GPU、训练数据集的大小等因素的影响。
在 PyTorch 中，torch.utils.data.Dataset 和 torch.utils.data.DataLoader 通常用于加载数据集和生成批处理。然而从版本 1.11 开始，PyTorch 引入了 TorchData 库，它实现了一种不同的加载数据集的办法。
在本文中，咱们将比拟数据集比拟大的状况下这两两种办法是如何工作的。咱们以 CelebA 和 DigiFace1M 的面部图像为例。表 1 显示了它们的比拟特色。咱们训练应用 ResNet-50 模型。而后进行 1 轮的训练来进行应用办法和工夫的比拟。
数据集的信息如下：
CelebA (align) 图片数：202,599 总大小：1.4 图片大小：178×218
DigiFace1M 图片数：720,000 总大小：14.6 图片大小：112×112
咱们应用的环境如下：
CPU: Intel(R) Core(TM) i9-9900K CPU @ 3.60GHz(16 核)
GPU: GeForce RTX 2080 Ti 11Gb
驱动版本 515.65.01 / CUDA 11.7 / CUDNN 8.4.0.27
Docker 20.10.21
Pytorch 1.12.1
TrochData 0.4.1
训练的代码如下：
| def train(data_loader: torch.utils.data.DataLoader, cfg: Config): | |
| # create model | |
| model = resnet50(num_classes=cfg.n_celeba_classes + cfg.n_digiface1m_classes, pretrained=True) | |
| torch.cuda.set_device(cfg.gpu) | |
| model = model.cuda(cfg.gpu) | |
| model.train() | |
| # define loss function (criterion) and optimizer | |
| criterion = torch.nn.CrossEntropyLoss().cuda(cfg.gpu) | |
| optimizer = torch.optim.SGD(model.parameters(), lr=0.1, | |
| momentum=0.9, | |
| weight_decay=1e-4) | |
| start_time = time.time() | |
| for _ in range(cfg.epochs): | |
| scaler = torch.cuda.amp.GradScaler(enabled=cfg.use_amp) | |
| for batch_idx, (images, target) in enumerate(data_loader): | |
| images = images.cuda(cfg.gpu, non_blocking=True) | |
| target = target.cuda(cfg.gpu, non_blocking=True) | |
| # compute output | |
| with torch.cuda.amp.autocast(enabled=cfg.use_amp): | |
| output = model(images) | |
| loss = criterion(output, target) | |
| # compute gradient | |
| scaler.scale(loss).backward() | |
| # do SGD step | |
| scaler.step(optimizer) | |
| scaler.update() | |
| optimizer.zero_grad() | |
| print(batch_idx, loss.item()) | |
| print(f'{time.time() - start_time} sec') |
Dataset
首先看看 Dataset, 这是自从 Pytorch 公布以来始终应用的形式，咱们对这个应该十分相熟。PyTorch 反对两种类型的数据集：map-style Datasets 和 iterable-style Datasets。Map-style Dataset 在事后晓得元素个数的状况下应用起来很不便。该类实现了__getitem__()和__len__()办法。如果通过索引读取太费时间或者无奈取得，那么能够应用 iterable-style，须要实现__iter__() 办法。在咱们的例子中，map-style 曾经能够了，因为对于 CelebA 和 DigiFace1M 数据集，咱们晓得其中的图像总数。
上面咱们创立 CelebADataset 类。对于 CelebA，类标签位于 identity_CelebA.txt 文件中。CelebA 和 DigiFace1M 中的面部图像在裁剪方面有所不同，因而为了在图像上传后缩小 getitem 办法中的这些差别，必须从各个方面略微裁剪它们。
| from PIL import Image | |
| from torch.utils.data import Dataset | |
| class CelebADataset(torch.utils.data.Dataset): | |
| def __init__(self, data_path: str, transform) -> None: | |
| self.data_path = data_path | |
| self.transform = transform | |
| self.image_names, self.labels = self.load_labels(f'{data_path}/identity_CelebA.txt') | |
| def __len__(self) -> int: | |
| return len(self.image_names) | |
| def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]: | |
| image_path = f'{self.data_path}/img_align_celeba/{self.image_names[idx]}' | |
| image = Image.open(image_path) | |
| left, right, top, bottom = 25, 153, 45, 173 | |
| image = image.crop((left, top, right, bottom)) | |
| if self.transform is not None: | |
| image = self.transform(image) | |
| label = self.labels[idx] | |
| return image, label | |
| @staticmethod | |
| def load_labels(labels_path: str) -> Tuple[list, list]: | |
| image_names, labels = [], [] | |
| with open(labels_path, 'r', encoding='utf-8') as labels_file: | |
| lines = labels_file.readlines() | |
| for line in lines: | |
| file_name, class_id = line.split(' ') | |
| image_names.append(file_name) | |
| labels.append(int(class_id[:-1])) | |
| return image_names, labels |
对于 DigiFace1M 数据集，同一类的所有图像都在一个独自的文件夹中。然而这两个数据集中，类的标签是雷同的，所以对于在 DigiFace1M 咱们不须要获取类别，而是在 CelebA 中按类减少。所以咱们须要 add_to_class 变量。另外就是 DigiFace1M 中的图像以“RGBA”格局存储，因而仍需将其转换为“RGB”。
| class DigiFace1M(torch.utils.data.Dataset): | |
| def __init__(self, data_path: str, transform, add_to_class: int = 0) -> None: | |
| self.data_path = data_path | |
| self.transform = transform | |
| self.image_paths, self.labels = self.load_labels(data_path, add_to_class) | |
| def __len__(self): | |
| return len(self.image_paths) | |
| def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]: | |
| image = Image.open(self.image_paths[idx]).convert('RGB') | |
| if self.transform is not None: | |
| image = self.transform(image) | |
| label = self.labels[idx] | |
| return image, label | |
| @staticmethod | |
| def load_labels(data_path: str, add_to_class: int) -> Tuple[list, list]: | |
| image_paths, labels = [], [] | |
| for root, _, files in os.walk(data_path): | |
| for file_name in files: | |
| if file_name.endswith('.png'): | |
| image_paths.append(f'{root}/{file_name}') | |
| labels.append(int(os.path.basename(root)) + add_to_class) | |
| return image_paths, labels |
当初咱们能够应用 torch.utils.data 将两个数据汇合并为一个数据集 ConcatDataset，创立 DataLoader，开始训练。
| def main(): | |
| cfg = Config() | |
| celeba_dataset = CelebADataset(f'{cfg.data_path}/CelebA', cfg.transform) | |
| digiface_dataset = DigiFace1M(f'{cfg.data_path}/DigiFace1M', cfg.transform, cfg.n_celeba_classes) | |
| dataset = torch.utils.data.ConcatDataset([celeba_dataset, digiface_dataset]) | |
| loader = torch.utils.data.DataLoader( | |
| dataset=dataset, | |
| batch_size=cfg.batch_size, | |
| shuffle=True, | |
| drop_last=True, | |
| num_workers=cfg.n_workers) | |
| utils.train(loader, cfg) |
TorchData API
与 Dataset 一样，TorchData 反对 map-style 和 iterable-style 的数据处理管道。然而官网倡议应用 IterDataPipe，只在必要时将其转换为 MapDataPipe。
因为 TorchData 提供了优化的数据加载实用程序，能够帮忙咱们不便的构建解决流程。以下是一些次要的性能:
- IterableWrapper：包装可迭代对象以创立 IterDataPipe。
- FileListerr：给定目录的门路，将生成根目录内文件的文件路径名(path + filename)
- Filterr：依据输出 filter_fn(函数名:filter)从源数据口过滤元素
- Mapperr：对源 DataPipe 中的每个项利用函数(函数名:map)
- Concaterr：连贯多个可迭代数据管道(函数名:concat)
- Shufflerr：打乱输出 DataPipe 数据的程序(函数名:shuffle)
- ShardingFilterr：容许对 DataPipe 进行分片(函数名:sharding_filter)
应用 TorchData 构建 CelebA 和 DigiFace1M 的数据处理管道，咱们须要执行以下步骤:
对于 CelebA 数据集：创立一个列表(file_name, label，‘ CelebA ‘)，并应用 IterableWrapper 从它创立一个 IterDataPipe
对于 DigiFace1M：应用 FileLister 创立一个 IterDataPipe，返回所有图像文件的门路，应用 Mapper 来应用 collate_ann。这个函数以图像门路作为输出，并返回元组(file_name, label，‘ DigiFace1M ‘)。
下面两个步骤之后，咱们失去两个数据类型 (file_name, label, data_name) 的后果。而后应用 Concater 将它们连贯到一个数据管道中。
应用 Shufflerr，打乱程序，这与在 DataLoader 中设置了 shuffle=True 是一样的。
应用 ShardingFilter 将数据管道宰割成片。每个 worker 将领有原始 DataPipe 元素的 n 个局部，其中 n 等于 worker 的数量。（多线程解决，DataLoader 中的 num_worker）
最初就是从磁盘读取图像
残缺代码如下：
| @torchdata.datapipes.functional_datapipe("load_image") | |
| class ImageLoader(torchdata.datapipes.iter.IterDataPipe): | |
| def __init__(self, source_datapipe, **kwargs) -> None: | |
| self.source_datapipe = source_datapipe | |
| self.transform = kwargs['transform'] | |
| def __iter__(self) -> Tuple[torch.Tensor, int]: | |
| for file_name, label, data_name in self.source_datapipe: | |
| image = Image.open(file_name) | |
| if data_name == 'DigiFace1M': | |
| image = image.convert('RGB') | |
| elif data_name == 'CelebA': | |
| left, right, top, bottom = 25, 153, 45, 173 | |
| image = image.crop((left, top, right, bottom)) | |
| if self.transform is not None: | |
| image = self.transform(image) | |
| yield image, label | |
| def collate_ann(file_path): | |
| label = int(os.path.basename(os.path.dirname(file_path))) + N_CELEBA_CLASSES | |
| data_name = os.path.basename(os.path.dirname(os.path.dirname(file_path))) | |
| return file_path, label, data_name | |
| def load_celeba_labels(labels_path: str) -> Dict[str, int]: | |
| labels = [] | |
| data_path = os.path.split(labels_path)[0] | |
| with open(labels_path, 'r', encoding='utf-8') as labels_file: | |
| lines = labels_file.readlines() | |
| for line in lines: | |
| file_name, class_id = line.split(' ') | |
| class_id = int(class_id[:-1]) | |
| labels.append((f'{data_path}/img_align_celeba/{file_name}', class_id, 'CelebA')) | |
| return labels | |
| def build_datapipes(cfg: Config) -> torchdata.datapipes.iter.IterDataPipe: | |
| celeba_dp = torchdata.datapipes.iter.IterableWrapper( | |
| load_celeba_labels(labels_path=f'{cfg.data_path}/CelebA/identity_CelebA.txt')) | |
| digiface_dp = torchdata.datapipes.iter.FileLister(f'{cfg.data_path}/DigiFace1M', masks='*.png', recursive=True) | |
| digiface_dp = digiface_dp.map(collate_ann) | |
| datapipe = celeba_dp.concat(digiface_dp) | |
| datapipe = datapipe.shuffle(buffer_size=100000) | |
| datapipe = datapipe.sharding_filter() | |
| datapipe = datapipe.load_image(transform=cfg.transform) | |
| return datapipe |
Torch 的 DataLoader 是同时反对 Datasets 和 DataPipe 的，所以咱们能够间接应用
| def main(): | |
| cfg = Config() | |
| datapipe = build_datapipes(cfg) | |
| loader = torch.utils.data.DataLoader( | |
| dataset=datapipe, | |
| batch_size=cfg.batch_size, | |
| shuffle=True, | |
| drop_last=True, | |
| num_workers=cfg.n_workers) | |
| utils.train(loader, cfg) |
减速数据读取的一个小技巧
批处理中耗时最长的操作之一是从磁盘读取图片。为了缩小这个操作所破费的工夫，能够加载所有图像并将它们宰割成小的数据集，例如 10,000 张图像保留为.pickle 文件。在读取时每一个 worker 只有读取一个相应的 pickle 文件即可
| def prepare_data(): | |
| cfg = Config() | |
| cfg.transform = None | |
| os.makedirs(cfg.prepared_data_path, exist_ok=True) | |
| celeba_dataset = dataset_example.CelebADataset(f'{cfg.data_path}/CelebA', cfg.transform) | |
| digiface_dataset = dataset_example.DigiFace1M(f'{cfg.data_path}/DigiFace1M', cfg.transform, cfg.n_celeba_classes) | |
| dataset = torch.utils.data.ConcatDataset([celeba_dataset, digiface_dataset]) | |
| shard_size = 10000 | |
| next_shard = 0 | |
| data = [] | |
| shuffled_idxs = np.arange(len(dataset)) | |
| np.random.shuffle(shuffled_idxs) | |
| for idx in tqdm(shuffled_idxs): | |
| data.append(dataset[idx]) | |
| if len(data) == shard_size: | |
| with open(f'{cfg.prepared_data_path}/{next_shard}_shard.pickle', 'wb') as _file: | |
| pickle.dump(data, _file) | |
| next_shard += 1 | |
| data = [] | |
| with open(f'{cfg.prepared_data_path}/{next_shard}_shard.pickle', 'wb') as _file: | |
| pickle.dump(data, _file) |
上面就是应用 FileLister 收集.pickle 数据集的所有门路，按 worker 划分并在每个 worker 上加载.pickle 数据。
| @torchdata.datapipes.functional_datapipe("load_pickle_data") | |
| class PickleDataLoader(torchdata.datapipes.iter.IterDataPipe): | |
| def __init__(self, source_datapipe, **kwargs) -> None: | |
| self.source_datapipe = source_datapipe | |
| self.transform = kwargs['transform'] | |
| def __iter__(self) -> Tuple[torch.Tensor, int]: | |
| for file_name in self.source_datapipe: | |
| with open(file_name, 'rb') as _file: | |
| pickle_data = pickle.load(_file) | |
| for image, label in pickle_data: | |
| image = self.transform(image) | |
| yield image, label | |
| def build_datapipes(cfg: Config) -> torchdata.datapipes.iter.IterDataPipe: | |
| datapipe = torchdata.datapipes.iter.FileLister(cfg.prepared_data_path, masks='*.pickle') | |
| datapipe = datapipe.shuffle() | |
| datapipe = datapipe.sharding_filter() | |
| datapipe = datapipe.load_pickle_data(transform=cfg.transform) | |
| return datapipe |
数据加载比照
咱们比拟三种不同数据加载办法。对于所有测试，batch_size = 600。
n workersDatasets, secDataPipes, secDataPipe + pickle, sec10 3581 7986 7585 10034 2993 760
当在未筹备好的数据上应用 DataPipe 进行训练时(不应用 pickle)，前几百个批次生成十分快，GPU 使用率简直是 100%，但随后速度逐步降落，这种办法甚至比应用 n_workers=10 的数据集还要慢。尽管我了解这两种办法的速度是一样的因为执行的操作是一样的，但实际上却不一样
DataLoader 的最佳 n_workers 没有一个固定值，因为这取决于工作 (图像大小，图像预处理的复杂性等等) 和计算机配置(HDD vs SSD)。
当在有大量小图像的数据集上训练时，做数据的筹备是必要的的，比方将小文件组合成几个大文件，这样能够缩小从磁盘读取数据的工夫。然而应用这种办法须要在将数据写入 shard 之前彻底打乱数据，来防止学习收敛性好转。还须要抉择正当的 shard 大小（它应该足够大以避免磁盘问题并且足够小以无效地应用 datappipes 中的 Shuffler 打乱数据）。
最初本文的代码在这里，有趣味的能够自行测试比拟：
https://avoid.overfit.cn/post/d431289d4723430b882e189008aeb959
作者：Karina Ovchinnikova
