1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
| from torch.utils.data import IterableDataset
|
|
| def default_fn(data):
| return data
|
|
| class MapperIterDataPipe(IterableDataset):
|
| def __init__(self, datapipe, fn=default_fn):
| self.datapipe = datapipe
| self.fn = fn
|
| def set_epoch(self, epoch):
| self.datapipe.set_epoch(epoch)
|
| def __iter__(self):
| assert callable(self.fn)
| for data in self.datapipe:
| yield self.fn(data)
|
|
