RabbitRUI/ruispace

1import numbers2import os3import queue as Queue4import threading5 6import mxnet as mx7import numpy as np8import torch9from torch.utils.data import DataLoader, Dataset10from torchvision import transforms11 12 13class BackgroundGenerator(threading.Thread):14    def __init__(self, generator, local_rank, max_prefetch=6):15        super(BackgroundGenerator, self).__init__()16        self.queue = Queue.Queue(max_prefetch)17        self.generator = generator18        self.local_rank = local_rank19        self.daemon = True20        self.start()21 22    def run(self):23        torch.cuda.set_device(self.local_rank)24        for item in self.generator:25            self.queue.put(item)26        self.queue.put(None)27 28    def next(self):29        next_item = self.queue.get()30        if next_item is None:31            raise StopIteration32        return next_item33 34    def __next__(self):35        return self.next()36 37    def __iter__(self):38        return self39 40 41class DataLoaderX(DataLoader):42 43    def __init__(self, local_rank, **kwargs):44        super(DataLoaderX, self).__init__(**kwargs)45        self.stream = torch.cuda.Stream(local_rank)46        self.local_rank = local_rank47 48    def __iter__(self):49        self.iter = super(DataLoaderX, self).__iter__()50        self.iter = BackgroundGenerator(self.iter, self.local_rank)51        self.preload()52        return self53 54    def preload(self):55        self.batch = next(self.iter, None)56        if self.batch is None:57            return None58        with torch.cuda.stream(self.stream):59            for k in range(len(self.batch)):60                self.batch[k] = self.batch[k].to(device=self.local_rank, non_blocking=True)61 62    def __next__(self):63        torch.cuda.current_stream().wait_stream(self.stream)64        batch = self.batch65        if batch is None:66            raise StopIteration67        self.preload()68        return batch69 70 71class MXFaceDataset(Dataset):72    def __init__(self, root_dir, local_rank):73        super(MXFaceDataset, self).__init__()74        self.transform = transforms.Compose(75            [transforms.ToPILImage(),76             transforms.RandomHorizontalFlip(),77             transforms.ToTensor(),78             transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),79             ])80        self.root_dir = root_dir81        self.local_rank = local_rank82        path_imgrec = os.path.join(root_dir, 'train.rec')83        path_imgidx = os.path.join(root_dir, 'train.idx')84        self.imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r')85        s = self.imgrec.read_idx(0)86        header, _ = mx.recordio.unpack(s)87        if header.flag > 0:88            self.header0 = (int(header.label[0]), int(header.label[1]))89            self.imgidx = np.array(range(1, int(header.label[0])))90        else:91            self.imgidx = np.array(list(self.imgrec.keys))92 93    def __getitem__(self, index):94        idx = self.imgidx[index]95        s = self.imgrec.read_idx(idx)96        header, img = mx.recordio.unpack(s)97        label = header.label98        if not isinstance(label, numbers.Number):99            label = label[0]100        label = torch.tensor(label, dtype=torch.long)101        sample = mx.image.imdecode(img).asnumpy()102        if self.transform is not None:103            sample = self.transform(sample)104        return sample, label105 106    def __len__(self):107        return len(self.imgidx)108 109 110class SyntheticDataset(Dataset):111    def __init__(self, local_rank):112        super(SyntheticDataset, self).__init__()113        img = np.random.randint(0, 255, size=(112, 112, 3), dtype=np.int32)114        img = np.transpose(img, (2, 0, 1))115        img = torch.from_numpy(img).squeeze(0).float()116        img = ((img / 255) - 0.5) / 0.5117        self.img = img118        self.label = 1119 120    def __getitem__(self, index):121        return self.img, self.label122 123    def __len__(self):124        return 1000000125