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
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363 | @register_vector_type(tf.Tensor, EagerTensor, tf.IndexedSlices)
class TensorflowVector(Vector):
"""Vector subclass to store tensorflow tensors.
This Vector is designed to store a collection of named TensorFlow
tensors, enabling computations that are either applied to each and
every coefficient, or imply two sets of aligned coefficients (i.e.
two TensorflowVector with similar specifications).
Note that support for IndexedSlices is implemented, as these are a
common type for auto-differentiated gradients. When using built-in
operators and methods, these structures will be preserved, unless
densification is required (e.g. when summing with a dense tensor).
When using `TensorflowVector.apply_func` directly, support for the
IndexedSlices' preservation should be added manually, typically by
using `declearn.model.tensorflow.utils.add_indexed_slices_support`.
Note that this class does not currently support special tensor types
such as SparseTensor or RaggedTensor.
Use `vector.coefs` to access the stored coefficients.
Notes
-----
- A `TensorflowVector` can be operated with either a:
- scalar value
- `NumpyVector` that has similar specifications
- `TensorflowVector` that has similar specifications
- => resulting in a `TensorflowVector` in each of these cases.
- The wrapped tensors may be placed on any device (CPU, GPU...)
and may not be all on the same device.
- The device-placement of the initial `TensorflowVector`'s data
is preserved by operations, including with `NumpyVector`.
- When combining two `TensorflowVector`, the device-placement
of the left-most one is used; in that case, one ends up with
`gpu + cpu = gpu` while `cpu + gpu = cpu`. In both cases, a
warning will be emitted to prevent silent un-optimized copies.
- When deserializing a `TensorflowVector` (either by directly using
`TensorflowVector.unpack` or loading one from a JSON dump), loaded
tensors are placed based on the global device-placement policy
(accessed via `declearn.utils.get_device_policy`). Thus it may
have a different device-placement schema than at dump time but
should be coherent with that of `TensorflowModel` computations.
"""
@property
def _op_add(self) -> Callable[[Any, Any], Any]:
return tf_op_add
@property
def _op_sub(self) -> Callable[[Any, Any], Any]:
return tf_op_sub
@property
def _op_mul(self) -> Callable[[Any, Any], Any]:
return tf_op_mul
@property
def _op_div(self) -> Callable[[Any, Any], Any]:
return tf_op_div
@property
def _op_pow(self) -> Callable[[Any, Any], Any]:
return tf_op_pow
@property
def compatible_vector_types(self) -> Set[Type[Vector]]:
types = super().compatible_vector_types
return types.union({NumpyVector, TensorflowVector})
def __init__(
self, coefs: Dict[str, Union[tf.Tensor, tf.IndexedSlices]]
) -> None:
super().__init__(coefs)
def apply_func(
self,
func: Callable[..., Any],
*args: Any,
**kwargs: Any,
) -> Self:
if not getattr(func, "_pre_wrapped", False):
func = preserve_tensor_device(func) # pragma: no cover
return super().apply_func(func, *args, **kwargs)
def _apply_operation(
self,
other: Any,
func: Callable[[Any, Any], Any],
) -> Self:
if not getattr(func, "_pre_wrapped", False):
func = preserve_tensor_device(func) # pragma: no cover
return super()._apply_operation(other, func)
def shapes(
self,
) -> Dict[str, Tuple[int, ...]]:
return {
key: tuple(coef.shape.as_list())
for key, coef in self.coefs.items()
}
def dtypes(
self,
) -> Dict[str, str]:
return {key: coef.dtype.name for key, coef in self.coefs.items()}
def pack(
self,
) -> Dict[str, Any]:
data = {key: self._pack_tensor(tns) for key, tns in self.coefs.items()}
return data
@classmethod
def unpack(
cls,
data: Dict[str, Any],
) -> Self:
policy = get_device_policy()
device = select_device(gpu=policy.gpu, idx=policy.idx)
with tf.device(device):
coef = {key: cls._unpack_tensor(dat) for key, dat in data.items()}
return cls(coef)
@classmethod
def _pack_tensor(
cls,
tensor: Union[tf.Tensor, tf.IndexedSlices],
) -> Any:
"""Convert a Tensor to a JSON-serializable object."""
if isinstance(tensor, tf.IndexedSlices):
val = cls._pack_tensor(tensor.values)
ind = cls._pack_tensor(tensor.indices)
shp = cls._pack_tensor(tensor.dense_shape)
return ["slices", val, ind, shp]
return np.array(tensor.numpy())
@classmethod
def _unpack_tensor(
cls,
data: Any,
) -> Union[tf.Tensor, tf.IndexedSlices]:
"""Re-create a Tensor from a JSON-unpacked object."""
if isinstance(data, list) and (data[0] == "slices"):
val = cls._unpack_tensor(data[1])
ind = cls._unpack_tensor(data[2])
shp = cls._unpack_tensor(data[3])
return tf.IndexedSlices(val, ind, shp)
try:
return tf.convert_to_tensor(data)
except TypeError as exc: # pragma: no cover
raise TypeError("Invalid tf.Tensor dump received.") from exc
def __eq__(
self,
other: Any,
) -> bool:
valid = isinstance(other, TensorflowVector)
if valid:
valid = self.coefs.keys() == other.coefs.keys()
if valid:
valid = all(
self._tensor_equal(self.coefs[key], other.coefs[key])
for key in self.coefs
)
return valid
@staticmethod
def _tensor_equal(
t_a: Union[tf.Tensor, tf.IndexedSlices],
t_b: Union[tf.Tensor, tf.IndexedSlices],
) -> bool:
if not isinstance(t_a, type(t_b)):
return False
if isinstance(t_a, tf.IndexedSlices):
# fmt: off
return (
TensorflowVector._tensor_equal(t_a.indices, t_b.indices)
and TensorflowVector._tensor_equal(t_a.values, t_b.values)
)
with tf.device(t_a.device):
return tf.reduce_all(t_a == t_b).numpy()
def sign(self) -> Self:
return self.apply_func(tf_op_sign)
def minimum(
self,
other: Union[Self, float],
) -> Self:
if isinstance(other, Vector):
return self._apply_operation(other, tf_op_min)
return self.apply_func(tf_op_min, other)
def maximum(
self,
other: Union[Self, float],
) -> Self:
if isinstance(other, Vector):
return self._apply_operation(other, tf_op_max)
return self.apply_func(tf_op_max, other)
def sum(
self,
) -> Self:
return self.apply_func(tf.reduce_sum)
def __pow__(
self,
other: Any,
) -> Self:
# For square and square root, use dedicated functions rather
# than tf.pow as results tend to differ for small values.
if isinstance(other, (int, float)):
if other == 2:
return self.apply_func(tf_op_sqre)
if other == 0.5:
return self.apply_func(tf_op_sqrt)
return super().__pow__(other)
def get_vector_specs(
self,
) -> VectorSpec:
# Add IndexedSlices information to the base specs.
specs = super().get_vector_specs()
slices = [
name
for name in specs.names
if isinstance(self.coefs[name], tf.IndexedSlices)
]
if slices:
specs.kwargs["slices"] = slices
return specs
def flatten(
self,
) -> Tuple[List[float], VectorSpec]:
v_spec = self.get_vector_specs()
arrays = [] # type: List[np.ndarray]
for name in v_spec.names:
if isinstance(self.coefs[name], tf.IndexedSlices):
warnings.warn(
"Flattening a 'tf.IndexedSlices' structure into an "
"exhaustive list of floats. This may result in a high "
"memory cost.",
UserWarning,
)
arrays.append(
np.array(tf.convert_to_tensor(self.coefs[name]).numpy())
)
else:
arrays.append(np.array(self.coefs[name].numpy()))
values = flatten_numpy_arrays(arrays)
return values, v_spec
@classmethod
def unflatten(
cls,
values: List[float],
v_spec: VectorSpec,
) -> Self:
# Convert values to numpy arrays.
shapes = [v_spec.shapes[name] for name in v_spec.names]
dtypes = [v_spec.dtypes[name] for name in v_spec.names]
arrays = unflatten_numpy_arrays(values, shapes, dtypes)
# Gather information about IndexedSlices and prepare a list of tensors.
slices = v_spec.kwargs.get("slices", [])
tf_dat = {} # Dict[str, Union[tf.Tensor, tf.IndexedSlices]]
# Convert numpy arrays back to tensorflow structures.
# Operate under the current device-placement policy.
policy = get_device_policy()
device = select_device(gpu=policy.gpu, idx=policy.idx)
with tf.device(device):
for name, array in zip(v_spec.names, arrays):
# Recover IndexedSlices structures from dense arrays.
if name in slices:
non_zero = np.any(
array != 0.0, axis=tuple(range(1, array.ndim))
)
ind = np.where(non_zero)[0].astype("int32")
tf_dat[name] = tf.IndexedSlices(
values=tf.convert_to_tensor(array[non_zero]),
indices=tf.convert_to_tensor(ind),
dense_shape=tf.convert_to_tensor(array.shape),
)
# Otherwise, merely convert arrays to tensors.
else:
tf_dat[name] = tf.convert_to_tensor(array)
return cls(tf_dat)
|