| 76 | self.random_erasing = None |
| 77 | |
| 78 | def __iter__(self): |
| 79 | stream = torch.cuda.Stream() |
| 80 | first = True |
| 81 | |
| 82 | for next_input, next_target in self.loader: |
| 83 | with torch.cuda.stream(stream): |
| 84 | next_input = next_input.cuda(non_blocking=True) |
| 85 | next_target = next_target.cuda(non_blocking=True) |
| 86 | if self.fp16: |
| 87 | next_input = next_input.half().sub_(self.mean).div_(self.std) |
| 88 | else: |
| 89 | next_input = next_input.float().sub_(self.mean).div_(self.std) |
| 90 | if self.random_erasing is not None: |
| 91 | next_input = self.random_erasing(next_input) |
| 92 | |
| 93 | if not first: |
| 94 | yield input, target |
| 95 | else: |
| 96 | first = False |
| 97 | |
| 98 | torch.cuda.current_stream().wait_stream(stream) |
| 99 | input = next_input |
| 100 | target = next_target |
| 101 | |
| 102 | yield input, target |
| 103 | |
| 104 | def __len__(self): |
| 105 | return len(self.loader) |