| 23 | |
| 24 | |
| 25 | class _DatasetReader(Reader): |
| 26 | def __init__(self, dataset, name, batch_size=1, enforce_batch_size=False): |
| 27 | """Don't call this directly. Instead, use dataset.reader()""" |
| 28 | Reader.__init__(self, dataset.content()) |
| 29 | self.dataset = dataset |
| 30 | self.name = name or (dataset.name + '_cursor') |
| 31 | self.batch_size = batch_size |
| 32 | self.enforce_batch_size = enforce_batch_size |
| 33 | self.cursor = None |
| 34 | |
| 35 | def setup_ex(self, init_net, exit_net): |
| 36 | if self.cursor is None: |
| 37 | self.cursor = init_net.CreateTreeCursor( |
| 38 | [], |
| 39 | init_net.NextScopedBlob(self.name), |
| 40 | fields=self.dataset.fields) |
| 41 | |
| 42 | def read(self, read_net): |
| 43 | assert self.cursor, 'setup not called.' |
| 44 | content = self.dataset.content() |
| 45 | with core.NameScope(read_net.NextName(self.name)): |
| 46 | fields = read_net.ReadNextBatch( |
| 47 | [self.cursor] + content.field_blobs(), |
| 48 | content.field_names(), |
| 49 | batch_size=self.batch_size, |
| 50 | enforce_batch_size=self.enforce_batch_size) |
| 51 | fields = core.output_to_list(fields) |
| 52 | return (read_net.IsEmpty([fields[0]]), fields) |
| 53 | |
| 54 | def reset(self, net): |
| 55 | net.ResetCursor([self.cursor], []) |
| 56 | |
| 57 | |
| 58 | class _DatasetRandomReader(Reader): |
no outgoing calls
no test coverage detected
searching dependent graphs…