Wrap a list to a torch Dataset. It produces elements of the list as data.
| 60 | |
| 61 | |
| 62 | class DatasetFromList(data.Dataset): |
| 63 | """ |
| 64 | Wrap a list to a torch Dataset. It produces elements of the list as data. |
| 65 | """ |
| 66 | |
| 67 | def __init__(self, lst: list, copy: bool = True, serialize: bool = True): |
| 68 | """ |
| 69 | Args: |
| 70 | lst (list): a list which contains elements to produce. |
| 71 | copy (bool): whether to deepcopy the element when producing it, |
| 72 | so that the result can be modified in place without affecting the |
| 73 | source in the list. |
| 74 | serialize (bool): whether to hold memory using serialized objects, when |
| 75 | enabled, data loader workers can use shared RAM from master |
| 76 | process instead of making a copy. |
| 77 | """ |
| 78 | self._lst = lst |
| 79 | self._copy = copy |
| 80 | self._serialize = serialize |
| 81 | |
| 82 | def _serialize(data): |
| 83 | buffer = pickle.dumps(data, protocol=-1) |
| 84 | return np.frombuffer(buffer, dtype=np.uint8) |
| 85 | |
| 86 | if self._serialize: |
| 87 | logger = logging.getLogger(__name__) |
| 88 | logger.info( |
| 89 | "Serializing {} elements to byte tensors and concatenating them all ...".format( |
| 90 | len(self._lst) |
| 91 | ) |
| 92 | ) |
| 93 | self._lst = [_serialize(x) for x in self._lst] |
| 94 | self._addr = np.asarray([len(x) for x in self._lst], dtype=np.int64) |
| 95 | self._addr = np.cumsum(self._addr) |
| 96 | self._lst = np.concatenate(self._lst) |
| 97 | logger.info("Serialized dataset takes {:.2f} MiB".format(len(self._lst) / 1024 ** 2)) |
| 98 | |
| 99 | def __len__(self): |
| 100 | if self._serialize: |
| 101 | return len(self._addr) |
| 102 | else: |
| 103 | return len(self._lst) |
| 104 | |
| 105 | def __getitem__(self, idx): |
| 106 | if self._serialize: |
| 107 | start_addr = 0 if idx == 0 else self._addr[idx - 1].item() |
| 108 | end_addr = self._addr[idx].item() |
| 109 | bytes = memoryview(self._lst[start_addr:end_addr]) |
| 110 | return pickle.loads(bytes) |
| 111 | elif self._copy: |
| 112 | return copy.deepcopy(self._lst[idx]) |
| 113 | else: |
| 114 | return self._lst[idx] |
| 115 | |
| 116 | |
| 117 | class ToIterableDataset(data.IterableDataset): |
no outgoing calls
no test coverage detected