-
Notifications
You must be signed in to change notification settings - Fork 34
/
numpyop.py
317 lines (273 loc) · 16.2 KB
/
numpyop.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
# Copyright 2019 The FastEstimator Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from typing import Any, Callable, Dict, Iterable, List, MutableMapping, Optional, Sequence, TypeVar, Union
import numpy as np
import tensorflow as tf
import torch
from torch.utils.data.dataloader import default_collate
from fastestimator.backend._to_tensor import to_tensor
from fastestimator.op.op import Op, get_inputs_by_op, write_outputs_by_op
from fastestimator.types import FilteredData
from fastestimator.util.traceability_util import traceable
from fastestimator.util.util import pad_batch
Tensor = TypeVar('Tensor', tf.Tensor, torch.Tensor, np.ndarray)
@traceable()
class NumpyOp(Op):
"""An Operator class which takes and returns numpy data.
These Operators are used in fe.Pipeline to perform data pre-processing / augmentation. They may also be used in
fe.Network to perform postprocessing on data.
Args:
inputs: Key(s) from which to retrieve data from the data dictionary.
outputs: Key(s) under which to write the outputs of this Op back to the data dictionary.
mode: What mode(s) to execute this Op in. For example, "train", "eval", "test", or "infer". To execute
regardless of mode, pass None. To execute in all modes except for a particular one, you can pass an argument
like "!infer" or "!train".
ds_id: What dataset id(s) to execute this Op in. To execute regardless of ds_id, pass None. To execute in all
ds_ids except for a particular one, you can pass an argument like "!ds1".
"""
def __init__(self,
inputs: Union[None, str, Iterable[str]] = None,
outputs: Union[None, str, Iterable[str]] = None,
mode: Union[None, str, Iterable[str]] = None,
ds_id: Union[None, str, Iterable[str]] = None) -> None:
super().__init__(inputs=inputs, outputs=outputs, mode=mode, ds_id=ds_id)
# in_place_edits tracks whether the .forward() method of this op will perform in-place edits of numpy arrays.
# This is inferred automatically by the system and is used for memory management optimization. If you are
# developing a NumpyOp which does in-place edits, the best practice is to set this to True in your init method.
self.in_place_edits = False
def forward(self, data: Union[np.ndarray, List[np.ndarray]],
state: Dict[str, Any]) -> Union[None, FilteredData, np.ndarray, List[np.ndarray]]:
"""A method which will be invoked in order to transform data.
This method will be invoked on individual elements of data before any batching / axis expansion is performed.
Args:
data: The arrays from the data dictionary corresponding to whatever keys this Op declares as its `inputs`.
state: Information about the current execution context, for example {"mode": "train"}.
Returns:
The `data` after applying whatever transform this Op is responsible for. It will be written into the data
dictionary based on whatever keys this Op declares as its `outputs`.
"""
return data
def forward_batch(self, data: Union[np.ndarray, List[np.ndarray]],
state: Dict[str, Any]) -> Union[None, FilteredData, np.ndarray, List[np.ndarray]]:
"""A method which will be invoked in order to transform a batch of data.
This method will be invoked on batches of data during network postprocessing. It should expect to receive
batched data and should itself return batched data.
Args:
data: The arrays from the data dictionary corresponding to whatever keys this Op declares as its `inputs`.
state: Information about the current execution context, for example {"mode": "train"}.
Returns:
The `data` after applying whatever transform this Op is responsible for. It will be written into the data
dictionary based on whatever keys this Op declares as its `outputs`.
"""
if isinstance(data, list):
data = [elem for elem in map(list, zip(*data))]
else:
data = [elem for elem in data]
results = [self.forward(elem, state) for elem in data]
if self.out_list:
results = [np.array(col) for col in [[row[i] for row in results] for i in range(len(results[0]))]]
else:
results = np.array(results)
return results
@traceable()
class Batch(NumpyOp):
"""Convert data instances into a batch of data.
Only one instance of a Batch Op can be present for a given epoch/mode/ds_id combination. Any Ops after this one will
operate on batches of data rather than individual instances (using their batch_forward methods).
Args:
batch_size: The batch size to use. If set, this will override any value specified by the Pipeline, allowing
control of the batch size on a per-mode and per-ds_id level. Note that this value will be ignored when using
a BatchDataset (or any dataset which decides on its own batch configuration).
drop_last: Whether to drop the last batch if the last batch is incomplete. Note that setting this to True when
using a BatchDataset (or any dataset which decides on its own batch configuration) won't do anything.
pad_value: The padding value if batch padding is needed. None indicates that no padding is needed. Mutually
exclusive with `collate_fn`.
collate_fn: A function to merge a list of data elements into a batch of data. Mutually exclusive with
`pad_value`.
mode: What mode(s) to execute this Op in. For example, "train", "eval", "test", or "infer". To execute
regardless of mode, pass None. To execute in all modes except for a particular one, you can pass an argument
like "!infer" or "!train".
ds_id: What dataset id(s) to execute this Op in. To execute regardless of ds_id, pass None. To execute in all
ds_ids except for a particular one, you can pass an argument like "!ds1".
"""
collate_fn: Callable[[List[MutableMapping[str, Any]]], Dict[str, Any]]
def __init__(self,
batch_size: Optional[int] = None,
drop_last: bool = False,
pad_value: Optional[Union[int, float]] = None,
collate_fn: Optional[Callable[[List[Dict[str, Any]]], Dict[str, Any]]] = None,
mode: Union[None, str, Iterable[str]] = None,
ds_id: Union[None, str, Iterable[str]] = None) -> None:
super().__init__(mode=mode, ds_id=ds_id)
if batch_size is not None:
if not isinstance(batch_size, int):
raise ValueError(f"batch_size must be an integer, but got {type(batch_size)}")
if batch_size < 0:
raise ValueError("batch_size must be non-negative")
self.batch_size = batch_size
self.drop_last = drop_last
if pad_value is not None and collate_fn is not None:
raise ValueError("Provide either a pad_value or collate_fn, but not both")
self._pad_value = pad_value
self.collate_fn = collate_fn
if pad_value is not None:
self.collate_fn = self._pad_batch_collate
if self.collate_fn is None:
# Note that this might get ignored in favor of default_convert inside the FEDataLoader if it looks like the
# user really doesn't want stuff to be batched.
self.collate_fn = default_collate
def _pad_batch_collate(self, batch: List[MutableMapping[str, Any]]) -> Dict[str, Any]:
"""A collate function which pads a batch of data.
Args:
batch: The data to be batched and collated.
Returns:
A padded and collated batch of data.
"""
pad_batch(batch, self._pad_value)
return default_collate(batch)
@traceable()
class Delete(NumpyOp):
"""Delete key(s) and their associated values from the data dictionary.
The system has special logic to detect instances of this Op and delete its `inputs` from the data dictionary.
Args:
keys: Existing key(s) to be deleted from the data dictionary.
mode: What mode(s) to execute this Op in. For example, "train", "eval", "test", or "infer". To execute
regardless of mode, pass None. To execute in all modes except for a particular one, you can pass an argument
like "!infer" or "!train".
ds_id: What dataset id(s) to execute this Op in. To execute regardless of ds_id, pass None. To execute in all
ds_ids except for a particular one, you can pass an argument like "!ds1".
"""
def __init__(self,
keys: Union[str, Sequence[str]],
mode: Union[None, str, Iterable[str]] = None,
ds_id: Union[None, str, Iterable[str]] = None) -> None:
super().__init__(inputs=keys, mode=mode, ds_id=ds_id)
def forward(self, data: Union[np.ndarray, List[np.ndarray]], state: Dict[str, Any]) -> None:
pass
def forward_batch(self, data: Union[Tensor, List[Tensor]], state: Dict[str, Any]) -> None:
pass
@traceable()
class LambdaOp(NumpyOp):
"""An Operator that performs any specified function as forward function.
Args:
fn: The function to be executed.
inputs: Key(s) from which to retrieve data from the data dictionary.
outputs: Key(s) under which to write the outputs of this Op back to the data dictionary.
mode: What mode(s) to execute this Op in. For example, "train", "eval", "test", or "infer". To execute
regardless of mode, pass None. To execute in all modes except for a particular one, you can pass an argument
like "!infer" or "!train".
ds_id: What dataset id(s) to execute this Op in. To execute regardless of ds_id, pass None. To execute in all
ds_ids except for a particular one, you can pass an argument like "!ds1".
"""
def __init__(self,
fn: Callable,
inputs: Union[None, str, Iterable[str]] = None,
outputs: Union[None, str, Iterable[str]] = None,
mode: Union[None, str, Iterable[str]] = None,
ds_id: Union[None, str, Iterable[str]] = None) -> None:
super().__init__(inputs=inputs, outputs=outputs, mode=mode, ds_id=ds_id)
self.fn = fn
self.in_list = True
def set_rua_level(self, magnitude_coef: float) -> None:
"""A method which will be invoked by the RUA Op to adjust the augmentation intensity.
Args:
magnitude_coef: The desired augmentation intensity (range [0-1]).
"""
def forward(self, data: List[np.ndarray], state: Dict[str, Any]) -> Union[np.ndarray, List[np.ndarray]]:
return self.fn(*data)
def forward_batch(self, data: Union[Tensor, List[Tensor]], state: Dict[str,
Any]) -> Union[np.ndarray, List[np.ndarray]]:
return self.forward(data, state)
@traceable()
class RemoveIf(NumpyOp):
"""An Operator which will remove a datapoint from the pipeline if the given criterion is satisfied.
Args:
fn: A function taking any desired `inputs` and returning a boolean. If the return value is true, the current
datapoint (or batch if using a batch dataset) will be removed and another will take its place.
replacement: Whether to replace the filtered element with another (thus maintaining the number of steps in an
epoch but potentially increasing data repetition) or else shortening the epoch by the number of filtered
data points (fewer steps per epoch than expected, but no extra data repetition). Either way, the number of
data points within an individual batch will remain the same. Even if `replacement` is true, data will not be
repeated until all of the given epoch's data has been traversed (except for at most 1 batch of data which
might not appear until after the re-shuffle has occurred).
inputs: Key(s) from which to retrieve data from the data dictionary.
mode: What mode(s) to execute this Op in. For example, "train", "eval", "test", or "infer". To execute
regardless of mode, pass None. To execute in all modes except for a particular one, you can pass an argument
like "!infer" or "!train".
ds_id: What dataset id(s) to execute this Op in. To execute regardless of ds_id, pass None. To execute in all
ds_ids except for a particular one, you can pass an argument like "!ds1".
"""
def __init__(self,
fn: Callable[..., bool],
replacement: bool = True,
inputs: Union[None, str, Iterable[str]] = None,
mode: Union[None, str, Iterable[str]] = None,
ds_id: Union[None, str, Iterable[str]] = None) -> None:
super().__init__(inputs=inputs, mode=mode, ds_id=ds_id)
self.filter_fn = fn
self.in_list = True
self.replacement = replacement
def forward(self, data: List[np.ndarray], state: Dict[str, Any]) -> Optional[FilteredData]:
if self.filter_fn(*data):
return FilteredData(replacement=self.replacement)
return None
def forward_batch(self, data: Union[Tensor, List[Tensor]], state: Dict[str, Any]) -> Optional[FilteredData]:
return self.forward(data, state)
def forward_numpyop(ops: List[NumpyOp],
data: MutableMapping[str, Any],
state: Dict[str, Any],
batched: Optional[str] = None,
shared: bool = True) -> Optional[FilteredData]:
"""Call the forward function for list of NumpyOps, and modify the data dictionary in place.
Args:
ops: A list of NumpyOps to execute.
data: The data dictionary.
state: Information about the current execution context, ex. {"mode": "train"}. Must contain at least the mode.
batched: Whether the `data` is batched or not. If it is batched, provide the string ('tf', 'torch', or 'np')
indicating which type of tensors the batch contains.
shared: Whether you want to place the resulting data into multi-processing shared memory. Only applicable when
`batched` is 'torch'.
"""
if not ops:
# Shortcut to prevent wasting time in to_tensor calls if there aren't any ops
return None
if batched:
# Cast data to Numpy before performing batch forward
for key, val in data.items():
data[key] = to_tensor(val, target_type='np')
for op in ops:
op_data = get_inputs_by_op(op, data, copy_on_write=op.in_place_edits)
try:
op_data = op.forward_batch(op_data, state) if batched else op.forward(op_data, state)
except ValueError as err:
if err.args[0] == 'assignment destination is read-only':
# If the numpy error text changes we'll need to make adjustments in the future
op.in_place_edits = True
op_data = get_inputs_by_op(op, data, copy_on_write=op.in_place_edits)
op_data = op.forward_batch(op_data, state) if batched else op.forward(op_data, state)
else:
raise err
if isinstance(op_data, FilteredData):
return op_data
if isinstance(op, Delete):
for key in op.inputs:
del data[key]
if op.outputs:
write_outputs_by_op(op, data, op_data)
if batched:
# Cast data back to original tensor type after performing batch forward
for key, val in data.items():
data[key] = to_tensor(val, target_type=batched, shared_memory=shared)
return None