Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,695 Bytes
2d9a728 |
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 |
import torch
import torch.distributed as dist
from utils.distributed import get_rank, is_dist_avail_and_initialized, is_main_process
import random
import logging
logger = logging.getLogger(__name__)
class MetaLoader(object):
""" wraps multiple data loader """
def __init__(self, name2loader):
"""Iterates over multiple dataloaders, it ensures all processes
work on data from the same dataloader. This loader will end when
the shorter dataloader raises StopIteration exception.
loaders: Dict, {name: dataloader}
"""
self.name2loader = name2loader
self.name2iter = {name: iter(l) for name, l in name2loader.items()}
name2index = {name: idx for idx, (name, l) in enumerate(name2loader.items())}
index2name = {v: k for k, v in name2index.items()}
iter_order = []
for n, l in name2loader.items():
iter_order.extend([name2index[n]]*len(l))
random.shuffle(iter_order)
iter_order = torch.Tensor(iter_order).to(torch.device("cuda")).to(torch.uint8)
# sync
if is_dist_avail_and_initialized():
# make sure all processes have the same order so that
# each step they will have data from the same loader
dist.broadcast(iter_order, src=0)
self.iter_order = [index2name[int(e.item())] for e in iter_order.cpu()]
logger.info(str(self))
def __str__(self):
output = [f"MetaLoader has {len(self.name2loader)} dataloaders, {len(self)} batches in total"]
for idx, (name, loader) in enumerate(self.name2loader.items()):
output.append(
f"dataloader index={idx} name={name}, batch-size={loader.batch_size} length(#batches)={len(loader)} "
)
return "\n".join(output)
def __len__(self):
return len(self.iter_order)
def __iter__(self):
""" this iterator will run indefinitely """
for name in self.iter_order:
_iter = self.name2iter[name]
batch = next(_iter)
yield name, batch
class MetaLoader_rs(object):
""" wraps multiple data loader """
def __init__(self, name2loader, skip_num=0):
"""Iterates over multiple dataloaders, it ensures all processes
work on data from the same dataloader. This loader will end when
the shorter dataloader raises StopIteration exception.
loaders: Dict, {name: dataloader}
"""
self.name2loader = name2loader
name2index = {name: idx for idx, (name, l) in enumerate(name2loader.items())}
index2name = {v: k for k, v in name2index.items()}
iter_order = []
for n, l in name2loader.items():
iter_order.extend([name2index[n]]*len(l))
random.shuffle(iter_order)
iter_order = torch.Tensor(iter_order).to(torch.device("cuda")).to(torch.uint8)
# sync
if is_dist_avail_and_initialized():
# make sure all processes have the same order so that
# each step they will have data from the same loader
dist.broadcast(iter_order, src=0)
if skip_num > 0:
iter_order_skip = iter_order[:skip_num]
for k, v in index2name.items():
media_step = (iter_order_skip == k).sum().item()
name2loader[v].sampler.set_start_iter(media_step)
logger.info(f"{v} dataloder skip steps: {media_step}")
iter_order = iter_order[skip_num:]
self.name2loader = name2loader
else:
logger.info("Do not skip steps for any dataloader!")
for k, v in index2name.items():
name2loader[v].sampler.set_start_iter(0)
self.name2iter = {name: iter(l) for name, l in name2loader.items()}
self.iter_idx = iter_order
self.iter_order = [index2name[int(e.item())] for e in iter_order.cpu()]
logger.info(str(self))
def __str__(self):
output = [f"MetaLoader has {len(self.name2loader)} dataloaders, {len(self)} batches in total"]
for idx, (name, loader) in enumerate(self.name2loader.items()):
length = (self.iter_idx == idx).sum()
output.append(
f"dataloader index={idx} name={name}, batch-size={loader.batch_size} length(#batches)={length} "
)
return "\n".join(output)
def __len__(self):
return len(self.iter_order)
def __iter__(self):
""" this iterator will run indefinitely """
for name in self.iter_order:
_iter = self.name2iter[name]
batch = next(_iter)
yield name, batch |