| 56 | |
| 57 | |
| 58 | class _DatasetRandomReader(Reader): |
| 59 | def __init__(self, dataset, name, indices, batch_size=1, loop_over=False, |
| 60 | enforce_batch_size=False): |
| 61 | """Don't call this directly. Instead, use dataset.random_reader()""" |
| 62 | Reader.__init__(self, dataset.content()) |
| 63 | self.dataset = dataset |
| 64 | self.cursor = None |
| 65 | self.name = name or (dataset.name + '_cursor') |
| 66 | self.indices = indices |
| 67 | self.batch_size = batch_size |
| 68 | self.loop_over = loop_over |
| 69 | self.enforce_batch_size = enforce_batch_size |
| 70 | |
| 71 | def setup_ex(self, init_net, exit_net): |
| 72 | if self.cursor is None: |
| 73 | self.cursor = init_net.CreateTreeCursor( |
| 74 | [], |
| 75 | init_net.NextScopedBlob(self.name), |
| 76 | fields=self.dataset.fields) |
| 77 | |
| 78 | def reset(self, net): |
| 79 | net.ResetCursor([self.cursor], []) |
| 80 | |
| 81 | def computeoffset(self, net): |
| 82 | self.reset(net) |
| 83 | offsets = net.ComputeOffset( |
| 84 | [self.cursor] + self.dataset.content().field_blobs(), |
| 85 | 'offsets') |
| 86 | self.offsets = offsets |
| 87 | |
| 88 | def sort_and_shuffle(self, net, sort_by_field=None, |
| 89 | shuffle_size=1, batch_size=1): |
| 90 | # no sorting by default |
| 91 | content = self.dataset.content() |
| 92 | sort_by_field_idx = -1 |
| 93 | if sort_by_field: |
| 94 | assert sort_by_field in content.field_names(), ( |
| 95 | 'Must be valid field.') |
| 96 | sort_by_field_idx = content.field_names().index(sort_by_field) |
| 97 | self.reset(net) |
| 98 | |
| 99 | indices = net.SortAndShuffle( |
| 100 | [self.cursor] + content.field_blobs(), |
| 101 | 'indices', |
| 102 | sort_by_field_idx=sort_by_field_idx, |
| 103 | shuffle_size=shuffle_size, |
| 104 | batch_size=batch_size) |
| 105 | self.indices = indices |
| 106 | |
| 107 | def read(self, read_net): |
| 108 | assert self.cursor, 'setup_ex not called' |
| 109 | assert self.indices, 'sort_and_shuffle not called' |
| 110 | assert self.offsets, 'computeoffset not called' |
| 111 | content = self.dataset.content() |
| 112 | with core.NameScope(read_net.NextName(self.name)): |
| 113 | fields = read_net.ReadRandomBatch( |
| 114 | [self.cursor, self.indices, self.offsets] + ( |
| 115 | content.field_blobs()), |
no outgoing calls
no test coverage detected
searching dependent graphs…