File size: 4,274 Bytes
37b3db0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
import torch
import mmap
import os
import json
from typing import Any, Callable, Optional
import numpy as np
import random
import tarfile
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from io import BytesIO
from PIL import Image, ImageDraw, ImageFilter
from training.distributed import world_info_from_env
from open_clip import tokenize
from .data import DataInfo
class IterativeWebDataset(torch.utils.data.IterableDataset):
def __init__(self, args, transform, tokenize):
self.args = args
start, end = os.path.basename(args.train_data).split("{")[1].split("}")[0].split("..")
self.num_shards = int(end) - int(start)
self.root_dir = os.path.dirname(args.train_data)
self.transform = transform
self.tokenizer = tokenize
self.start_shard_id = 0
self.shard_ids = list(range(self.num_shards))
def set_epoch(self, epoch, num_batches, step=0):
random.seed(epoch+step)
self.shard_ids = list(range(self.num_shards))
random.shuffle(self.shard_ids)
self.start_shard_id = (num_batches * epoch) % self.num_shards
def _get_tarball_path(self, shard_id):
return os.path.join(self.root_dir, f"{shard_id % 100}", f"{shard_id}.tar")
def _get_next_shard_id(self, shard_id):
shard_id += self.group_size
return shard_id % self.num_shards
def __iter__(self):
worker_info = torch.utils.data.get_worker_info()
if worker_info is None:
num_workers = 1
worker_id = 0
else:
num_workers = worker_info.num_workers
worker_id = worker_info.id
_, global_rank, world_size = world_info_from_env()
self.group_size = int(num_workers * world_size)
shard_id = num_workers * global_rank + worker_id
shard_id = (shard_id + self.start_shard_id) % self.num_shards
shard_id = self.shard_ids[shard_id]
while True:
tarball_path = self._get_tarball_path(shard_id)
if not os.path.exists(tarball_path):
shard_id = self._get_next_shard_id(shard_id)
continue
with tarfile.open(tarball_path) as tar:
members = tar.getmembers()
# metaclip_v1 can be iterative but the paper uses mmap for random access.
json_uuid, img_uuid = -1, -2
for member in members:
if member.name.endswith(".json"):
json_uuid = member.name[:-len(".json")]
with tar.extractfile(member) as f:
text_json = json.load(f)
if member.name.endswith(".jpeg"):
img_uuid = member.name[:-len(".jpeg")]
with tar.extractfile(member) as f:
img = f.read()
if img_uuid != json_uuid:
# assume uuid is json even and img ord;
continue
txt = random.choice(text_json["texts"])[1]
txt = self.tokenizer([txt])[0]
with Image.open(BytesIO(img)) as img:
image = img.convert("RGB")
image = self.transform(image)
yield image, txt
shard_id = self._get_next_shard_id(shard_id)
def get_metaclip_iter_wds_dataset(args, preprocess_fn, is_train, epoch=0):
# borrowed from get_csv_dataset
input_filename = args.train_data if is_train else args.val_data
assert input_filename
dataset = IterativeWebDataset(
args,
preprocess_fn,
tokenize,
)
assert is_train
num_samples = args.train_num_samples
sampler = None
dataloader = torch.utils.data.DataLoader(
dataset, sampler=sampler,
batch_size=args.batch_size,
num_workers=args.workers,
pin_memory=False,
drop_last=True,
)
dataloader.num_samples = num_samples
dataloader.num_batches = int(num_samples / (args.batch_size * args.world_size))
return DataInfo(dataloader, sampler)
|