| 15 | import torch.nn.functional as F |
| 16 | |
| 17 | class IterLoader: |
| 18 | def __init__(self, loader, length=None): |
| 19 | self.loader = loader |
| 20 | self.length = length |
| 21 | self.iter = None |
| 22 | |
| 23 | def __len__(self): |
| 24 | if self.length is not None: |
| 25 | return self.length |
| 26 | |
| 27 | return len(self.loader) |
| 28 | |
| 29 | def new_epoch(self): |
| 30 | self.iter = iter(self.loader) |
| 31 | |
| 32 | def next(self): |
| 33 | try: |
| 34 | return next(self.iter) |
| 35 | except: |
| 36 | self.iter = iter(self.loader) |
| 37 | return next(self.iter) |
| 38 | |
| 39 | |
| 40 | import os.path as osp |