Webdataset 加速深度学习数据加载

# read sample from the given tar file
k = rng.randint(0, len(buf) - 1)
sample, buf[k] = buf[k], sample
# return sample here

|Name          |  Read(MB/s)| Write(MB/s)|
|--------------|------------|------------|
|  SEQ1M Q1 T1 |     484.345|     352.739|
|  SEQ1M Q8 T1 |     516.135|     425.775|
| RND4K Q32T16 |     311.884|     282.172|
| . IOPS       |   76143.543|   68889.666|
| . latency us |       6.709|       7.422|
| RND4K Q1 T1  |      42.470|     106.964|
| . IOPS       |   10368.547|   26114.252|
| . latency us |       0.096|       0.037|

|Name          |  Read(MB/s)| Write(MB/s)|
|--------------|------------|------------|
|  SEQ1M Q1 T1 |     179.546|      79.943|
|  SEQ1M Q8 T1 |     172.259|      79.952|
| RND4K Q32T16 |       2.335|       0.929|
| . IOPS       |     569.994|     226.896|
| . latency us |     823.919|    1837.834|
| RND4K Q1 T1  |       0.686|       1.094|
| . IOPS       |     167.532|     267.147|
| . latency us |       5.962|       3.734|

import os
import random
import datetime
from multiprocessing import Process
from torchvision import datasets
from torchvision.datasets import ImageNet
from torchvision.datasets.folder import ImageFolder
from webdataset import TarWriter
def make_wds_shards(pattern, num_shards, num_workers, samples, map_func, **kwargs):
    random.shuffle(samples)
    samples_per_shards = [samples[i::num_shards] for i in range(num_shards)]
    shard_ids = list(range(num_shards))
    processes = [
        Process(
            target=write_partial_samples,
            args=(
                pattern,
                shard_ids[i::num_workers],
                samples_per_shards[i::num_workers],
                map_func,
                kwargs
        for i in range(num_workers)]
    for p in processes:
        p.start()
    for p in processes:
        p.join()
def write_partial_samples(pattern, shard_ids, samples, map_func, kwargs):
    for shard_id, samples in zip(shard_ids, samples):
        write_samples_into_single_shard(pattern, shard_id, samples, map_func, kwargs)
def write_samples_into_single_shard(pattern, shard_id, samples, map_func, kwargs):
    fname = pattern % shard_id
    print(f"[{datetime.datetime.now()}] start to write samples to shard {fname}")
    stream = TarWriter(fname, **kwargs)
    size = 0
    for item in samples:
        size += stream.write(map_func(item))
    stream.close()
    print(f"[{datetime.datetime.now()}] complete to write samples to shard {fname}")
    return size
if __name__ == "__main__":
    root = "/gdata/ImageNet2012/train"
    items = []
    dataset = ImageFolder(root=root,  loader=lambda x:x)
    for i in range(len(dataset)):
        items.append(dataset[i])
    print(dataset[0],os.path.splitext(os.path.basename(dataset[0][0]))[0])
    def map_func(item):
        name, class_idx = item
        with open(os.path.join(name), "rb") as stream:
            image = stream.read()
        sample = {
            "__key__": os.path.splitext(os.path.basename(name))[0],
            "jpg": image,
            "cls": str(class_idx).encode("ascii")
        return sample
    make_wds_shards(
        pattern="/userhome/tars/imagenet-1k-%06d.tar",
        num_shards=256, # 设置分片数量
        num_workers=8, # 设置创建wds数据集的进程数
        samples=items,
        map_func=map_func,
    )

import os
import time
import torch
import webdataset as wds
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
def get_ori_loader(disk, num_workers):
    def read_bytes(path):
        with open(path, "rb") as f:
            return f.read()
    root = "/mnt/extend/imagenet/train" if disk == "hdd" else "/home/chenyaofo/webdataset-test/train"
    dataset = ImageFolder(root,  loader=read_bytes, is_valid_file=lambda x: True)
    dataloader = DataLoader(dataset, num_workers=num_workers, shuffle=True, batch_size=128)
    return dataloader
def get_wds_loader(disk, num_workers):
    url = "/mnt/extend/tars/imagenet-1k-{000000..000256}.tar" if disk == "hdd" else "/home/chenyaofo/webdataset-test/tars/imagenet-1k-{000000..000256}.tar"
    def my_decoder(key, value):
        if not key.endswith(".jpg"):
            return None
        assert isinstance(value, bytes)
        return value
    dataset = wds.WebDataset(url).shuffle(1000).decode(my_decoder)
    dataloader = DataLoader(dataset, num_workers=num_workers, batch_size=128)
    return dataloader
def run_test(loader, disk):
    N_stop = 30000 if disk == "hdd" else 300000
    start = time.perf_counter()
    total_batch_size = 0
    total_bytes = 0
    for items in loader:
        if isinstance(items, dict):
            batch_size = len(items['jpeg.cls'])
            n_bytes = sum(map(lambda x: len(x), items['jpeg.jpg']))
        else:
            batch_size = len(items[1])
            n_bytes = sum(map(lambda x: len(x), items[0]))
        total_batch_size += batch_size
        # print(total_batch_size)
        total_bytes += n_bytes
        if total_batch_size > N_stop:
            end = time.perf_counter()
            return total_batch_size, total_bytes, end-start
for disk in ["ssd", "hdd"]: