| 138 | |
| 139 | |
| 140 | class ErasePythonPipeline(Pipeline): |
| 141 | def __init__( |
| 142 | self, |
| 143 | function, |
| 144 | batch_size, |
| 145 | data_layout, |
| 146 | iterator, |
| 147 | anchor, |
| 148 | shape, |
| 149 | axis_names, |
| 150 | axes, |
| 151 | fill_value, |
| 152 | erase_func=erase_func, |
| 153 | num_threads=1, |
| 154 | device_id=0, |
| 155 | ): |
| 156 | super(ErasePythonPipeline, self).__init__( |
| 157 | batch_size, num_threads, device_id, exec_async=False, exec_pipelined=False |
| 158 | ) |
| 159 | self.iterator = iterator |
| 160 | self.inputs = ops.ExternalSource() |
| 161 | self.data_layout = data_layout |
| 162 | |
| 163 | if isinstance(fill_value, RandomDataIterator): |
| 164 | self.fill_value_iterator = fill_value |
| 165 | self.fill_value_inputs = ops.ExternalSource() |
| 166 | fill_value = None |
| 167 | function = partial(erase_func, anchor, shape, axis_names, axes, data_layout) |
| 168 | else: |
| 169 | self.fill_value_iterator = None |
| 170 | function = partial(erase_func, anchor, shape, axis_names, axes, data_layout, fill_value) |
| 171 | |
| 172 | self.erase = ops.PythonFunction(function=function, output_layouts=data_layout) |
| 173 | |
| 174 | def define_graph(self): |
| 175 | self.data = self.inputs() |
| 176 | if self.fill_value_iterator is not None: |
| 177 | self.fill_value_data = self.fill_value_inputs() |
| 178 | out = self.erase(self.fill_value_data, self.data) |
| 179 | else: |
| 180 | out = self.erase(self.data) |
| 181 | return out |
| 182 | |
| 183 | def iter_setup(self): |
| 184 | data = self.iterator.next() |
| 185 | self.feed_input(self.data, data) |
| 186 | if self.fill_value_iterator is not None: |
| 187 | fill_value_data = self.fill_value_iterator.next() |
| 188 | self.feed_input(self.fill_value_data, fill_value_data) |
| 189 | |
| 190 | |
| 191 | def check_operator_erase_vs_python( |
no outgoing calls