| import numbers | |
| import os | |
| import queue as Queue | |
| import threading | |
| import mxnet as mx | |
| import numpy as np | |
| import torch | |
| from torch.utils.data import DataLoader, Dataset | |
| from torchvision import transforms | |
| class BackgroundGenerator(threading.Thread): | |
| def __init__(self, generator, local_rank, max_prefetch=6): | |
| super(BackgroundGenerator, self).__init__() | |
| self.queue = Queue.Queue(max_prefetch) | |
| self.generator = generator | |
| self.local_rank = local_rank | |
| self.daemon = True | |
| self.start() | |
| def run(self): | |
| torch.cuda.set_device(self.local_rank) | |
| for item in self.generator: | |
| self.queue.put(item) | |
| self.queue.put(None) | |
| def next(self): | |
| next_item = self.queue.get() | |
| if next_item is None: | |
| raise StopIteration | |
| return next_item | |
| def __next__(self): | |
| return self.next() | |
| def __iter__(self): | |
| return self | |
| class DataLoaderX(DataLoader): | |
| def __init__(self, local_rank, **kwargs): | |
| super(DataLoaderX, self).__init__(**kwargs) | |
| self.stream = torch.cuda.Stream(local_rank) | |
| self.local_rank = local_rank | |
| def __iter__(self): | |
| self.iter = super(DataLoaderX, self).__iter__() | |
| self.iter = BackgroundGenerator(self.iter, self.local_rank) | |
| self.preload() | |
| return self | |
| def preload(self): | |
| self.batch = next(self.iter, None) | |
| if self.batch is None: | |
| return None | |
| with torch.cuda.stream(self.stream): | |
| for k in range(len(self.batch)): | |
| self.batch[k] = self.batch[k].to(device=self.local_rank, non_blocking=True) | |
| def __next__(self): | |
| torch.cuda.current_stream().wait_stream(self.stream) | |
| batch = self.batch | |
| if batch is None: | |
| raise StopIteration | |
| self.preload() | |
| return batch | |
| class MXFaceDataset(Dataset): | |
| def __init__(self, root_dir, local_rank): | |
| super(MXFaceDataset, self).__init__() | |
| self.transform = transforms.Compose( | |
| [transforms.ToPILImage(), | |
| transforms.RandomHorizontalFlip(), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), | |
| ]) | |
| self.root_dir = root_dir | |
| self.local_rank = local_rank | |
| path_imgrec = os.path.join(root_dir, 'train.rec') | |
| path_imgidx = os.path.join(root_dir, 'train.idx') | |
| self.imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') | |
| s = self.imgrec.read_idx(0) | |
| header, _ = mx.recordio.unpack(s) | |
| if header.flag > 0: | |
| self.header0 = (int(header.label[0]), int(header.label[1])) | |
| self.imgidx = np.array(range(1, int(header.label[0]))) | |
| else: | |
| self.imgidx = np.array(list(self.imgrec.keys)) | |
| def __getitem__(self, index): | |
| idx = self.imgidx[index] | |
| s = self.imgrec.read_idx(idx) | |
| header, img = mx.recordio.unpack(s) | |
| label = header.label | |
| if not isinstance(label, numbers.Number): | |
| label = label[0] | |
| label = torch.tensor(label, dtype=torch.long) | |
| sample = mx.image.imdecode(img).asnumpy() | |
| if self.transform is not None: | |
| sample = self.transform(sample) | |
| return sample, label | |
| def __len__(self): | |
| return len(self.imgidx) | |
| class SyntheticDataset(Dataset): | |
| def __init__(self, local_rank): | |
| super(SyntheticDataset, self).__init__() | |
| img = np.random.randint(0, 255, size=(112, 112, 3), dtype=np.int32) | |
| img = np.transpose(img, (2, 0, 1)) | |
| img = torch.from_numpy(img).squeeze(0).float() | |
| img = ((img / 255) - 0.5) / 0.5 | |
| self.img = img | |
| self.label = 1 | |
| def __getitem__(self, index): | |
| return self.img, self.label | |
| def __len__(self): | |
| return 1000000 | |