| 177 | |
| 178 | |
| 179 | class _PreprocessMultiProcessingDataLoaderIter(_MultiProcessingDataLoaderIter): |
| 180 | |
| 181 | def __init__(self, loader, preprocess: Preprocess=None, queue_max_size=2): |
| 182 | self._preprocess = preprocess |
| 183 | if self._preprocess is not None and self._preprocess.has_cpu_preprocess(): |
| 184 | org_create_fetcher = _DatasetKind.create_fetcher |
| 185 | _DatasetKind.create_fetcher = _create_fetcher_proxy(org_create_fetcher, self._preprocess) |
| 186 | |
| 187 | super(_PreprocessMultiProcessingDataLoaderIter, self).__init__(loader) |
| 188 | |
| 189 | if self._preprocess is not None and self._preprocess.has_cpu_preprocess(): |
| 190 | _DatasetKind.create_fetcher = org_create_fetcher |
| 191 | |
| 192 | # self._preprocess_out_queue = self._data_queue |
| 193 | |
| 194 | if self._preprocess is not None and self._preprocess.has_gpu_preprocess(): |
| 195 | self._preprocess_in_queue = self._data_queue |
| 196 | self._data_queue = queue.Queue(maxsize=queue_max_size) |
| 197 | self._gpu_preprocess_thread_done_event = threading.Event() |
| 198 | self._device_id = torch.cuda.current_device() |
| 199 | self._timeout = self._timeout if self._timeout > 0 else (12 * _utils.MP_STATUS_CHECK_INTERVAL) |
| 200 | self._stream = torch.cuda.Stream() |
| 201 | |
| 202 | self._gpu_preprocess_thread = threading.Thread(target=self._gpu_preprocess_loop) |
| 203 | self._gpu_preprocess_thread.daemon = True |
| 204 | self._gpu_preprocess_thread.start() |
| 205 | |
| 206 | |
| 207 | def _gpu_preprocess_loop(self): |
| 208 | torch.set_num_threads(1) |
| 209 | # print("start _gpu_preprocess_loop") |
| 210 | torch.cuda.set_device(self._device_id) |
| 211 | while not self._gpu_preprocess_thread_done_event.is_set(): |
| 212 | try: |
| 213 | r = self._preprocess_in_queue.get(timeout=_utils.MP_STATUS_CHECK_INTERVAL) |
| 214 | # print(f"r = {r}") |
| 215 | except queue.Empty: |
| 216 | # print(f"_preprocess_in_queue empty, timeout = {self._timeout}") |
| 217 | continue |
| 218 | idx, data = r |
| 219 | # print(f"************** _PreprocessMultiProcessingDataLoaderIter _gpu_preprocess_loop **************") |
| 220 | # print(f"_gpu_preprocess_loop data is {type(data)}") |
| 221 | if not isinstance(idx, _utils.worker._ResumeIteration) and not self._gpu_preprocess_thread_done_event.is_set() and not isinstance(data, _utils.ExceptionWrapper): |
| 222 | try: |
| 223 | if self._preprocess is not None and self._preprocess.has_gpu_preprocess(): |
| 224 | with torch.no_grad(), torch.cuda.stream(self._stream): |
| 225 | data = self._preprocess.gpu_preprocess(data, self._stream) |
| 226 | # print(f"data = {data}") |
| 227 | except Exception as e: |
| 228 | # print(e) |
| 229 | data = _utils.ExceptionWrapper(where=f"in _gpu_preprocess_loop thread for device {self._device_id}") |
| 230 | r = (idx, data) |
| 231 | while not self._gpu_preprocess_thread_done_event.is_set(): |
| 232 | try: |
| 233 | self._data_queue.put(r, timeout=_utils.MP_STATUS_CHECK_INTERVAL) |
| 234 | break |
| 235 | except queue.Full: |
| 236 | # print(f"_data_queue full, timeout = {self._timeout}") |