File size: 6,061 Bytes
391690f | 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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | """cc12m_data.py — CC12M-resident data path for dist_bed (pod rung).
Trains straight from the pixparse/cc12m-wds tars via a one-time offset
index (no extraction: 1.18TB stays as 2,176 tars). Row order = the
feature-bank concatenation order (features_0000..2175), so targets[idx]
indexing works exactly like the COCO path.
python tools/cc12m_data.py --build-index # one-time, parallel scan
"""
import glob
import json
import os
import tarfile
import numpy as np
import torch
from PIL import Image, ImageFile
from torch.utils.data import Dataset
ImageFile.LOAD_TRUNCATED_IMAGES = True
DATA_ROOT = os.environ.get("DIST_DATA_ROOT", "./data")
CC12M_TARS = os.environ.get("CC12M_TARS",
os.path.join(DATA_ROOT, "cc12m"))
CC12M_FEATS = os.environ.get("CC12M_FEATS",
os.path.join(DATA_ROOT, "cc12m_feats",
"clip_b16_laion2b"))
INDEX = os.path.join(CC12M_TARS, "jpg_offset_index.npz")
def _scan_shard(path):
"""-> (shard_name, {key: (offset, size)}) for .jpg members."""
out = {}
with tarfile.open(path, "r") as tar:
for m in tar:
if m.isfile() and m.name.endswith(".jpg"):
out[m.name[:-4]] = (m.offset_data, m.size)
return os.path.basename(path), out
def build_index(workers=32):
from multiprocessing import Pool
tars = sorted(glob.glob(os.path.join(CC12M_TARS, "*.tar")))
assert tars, f"no tars under {CC12M_TARS}"
print(f"[index] scanning {len(tars)} tars with {workers} workers")
shard_names, keys, shard_idx, offs, sizes = [], [], [], [], []
with Pool(workers) as pool:
for name, d in pool.imap(_scan_shard, tars, chunksize=4):
si = len(shard_names)
shard_names.append(name)
for k, (o, s) in d.items():
keys.append(k)
shard_idx.append(si)
offs.append(o)
sizes.append(s)
if si % 200 == 0:
print(f" [{si}] {sum(map(len, [keys]))} keys", flush=True)
np.savez(INDEX, shards=np.array(shard_names),
keys=np.array(keys), shard_idx=np.array(shard_idx, np.int32),
off=np.array(offs, np.int64), size=np.array(sizes, np.int64))
print(f"[index] {len(keys)} jpgs -> {INDEX}")
def load_cc12m_tower(cfg):
"""-> (keys, fp16 UNNORMALIZED) for an extra tower bank
(cc12m_feats_{cfg} under CC12M_FEATS_ROOT), concat-cached."""
root = os.path.join(os.environ.get("CC12M_FEATS_ROOT", "/data"),
f"cc12m_feats_{cfg}")
cache = os.path.join(root, "all_concat.pt")
if os.path.isfile(cache):
d = torch.load(cache, map_location="cpu", weights_only=True)
return d["keys"], d["emb"]
files = sorted(glob.glob(os.path.join(root, "features_*.pt")))
assert len(files) == 2176, f"{cfg} bank incomplete: {len(files)}/2176"
keys, embs = [], []
for f in files:
try:
d = torch.load(f, map_location="cpu", weights_only=True)
except Exception:
# pod tower shards carry numpy.str_ keys (offset-index array
# slices) which weights_only refuses; these are SELF-AUTHORED
# files on our own disk — trusted load, keys coerced below.
d = torch.load(f, map_location="cpu", weights_only=False)
keys += [str(k) for k in d["keys"]]
embs.append(d["emb"])
emb = torch.cat(embs)
torch.save({"keys": keys, "emb": emb}, cache)
return keys, emb
def load_cc12m_gpa():
"""-> (keys, fp16 GPA targets) built by cc12m_gpa.py."""
p = os.path.join(os.environ.get("CC12M_FEATS_ROOT", "/data"),
"cc12m_gpa.pt")
d = torch.load(p, map_location="cpu", weights_only=True)
return d["keys"], d["emb"]
def load_cc12m_targets():
"""-> (keys list, fp16 targets (N,512) UNNORMALIZED) in bank file order.
Caller normalizes (the bed normalizes per-batch)."""
cache = os.path.join(os.path.dirname(CC12M_FEATS), "all_concat.pt")
if os.path.isfile(cache):
d = torch.load(cache, map_location="cpu", weights_only=True)
return d["keys"], d["emb"]
files = sorted(glob.glob(os.path.join(CC12M_FEATS, "features_*.pt")))
assert len(files) == 2176, f"bank incomplete: {len(files)}/2176"
keys, embs = [], []
for f in files:
d = torch.load(f, map_location="cpu", weights_only=True)
keys += d["keys"]
embs.append(d["emb"])
emb = torch.cat(embs)
torch.save({"keys": keys, "emb": emb}, cache)
return keys, emb
class Cc12mImages(Dataset):
"""(row i) -> (STUDENT_TF tensor, i), rows in bank order. Per-worker
persistent file handles; reads jpg bytes by offset from the tars."""
def __init__(self, keys, tf):
z = np.load(INDEX, allow_pickle=False)
pos = {k: i for i, k in enumerate(z["keys"].tolist())}
sel = np.array([pos[k] for k in keys], dtype=np.int64)
self.shards = [os.path.join(CC12M_TARS, s)
for s in z["shards"].tolist()]
self.shard_idx = z["shard_idx"][sel]
self.off = z["off"][sel]
self.size = z["size"][sel]
self.tf = tf
def __len__(self):
return len(self.off)
def _read(self, i):
# open/seek/read/close per item: with 2,176 shards x N workers,
# persistent handles exhaust the fd limit (learned the hard way,
# wave 1). Syscall cost is noise next to jpg decode.
with open(self.shards[int(self.shard_idx[i])], "rb") as fh:
fh.seek(int(self.off[i]))
return fh.read(int(self.size[i]))
def __getitem__(self, i):
import io
im = Image.open(io.BytesIO(self._read(i))).convert("RGB")
return self.tf(im), i
if __name__ == "__main__":
import argparse
ap = argparse.ArgumentParser()
ap.add_argument("--build-index", action="store_true")
ap.add_argument("--workers", type=int, default=32)
a = ap.parse_args()
if a.build_index:
build_index(a.workers)
|