Added TensorFlow support to nncf.Tensor
#3106
Conversation
olegkkruglov
force-pushed
the
tf-tensor-numeric
branch
from
44d27af to
79ded03
Compare
github-actions
bot
added
the
documentation
olegkkruglov
force-pushed
the
tf-tensor-numeric
branch
from
79ded03 to
c977245
Compare
olegkkruglov
force-pushed
the
tf-tensor-numeric
branch
from
c977245 to
9b6bf86
Compare
|
Please, update your branch from develop. |
olegkkruglov
force-pushed
the
tf-tensor-numeric
branch
from
9b6bf86 to
298891b
Compare
nncf/tensor/functions/tf_numeric.py
Outdated
| keepdims: bool = False, | ||
| ) -> tf.Tensor: | ||
| numpy_a = np.array(a) | ||
| numpy_median = np.median(numpy_a, axis=axis, keepdims=keepdims) |
There was a problem hiding this comment.
I would suggest to implement median using tf.math.top_k. I hope the better performance can be achieved on GPU using tf.math.top_k.
There was a problem hiding this comment.
Do you have any evidence that this implementation is more performant than calling np.median?
olegkkruglov
force-pushed
the
tf-tensor-numeric
branch
from
c687e6a to
bd17d4a
Compare
|
It looks like the PR is stuck. I see that @olegkkruglov answered all the comments from @alexsu52 |
| axis = (0, 1) | ||
|
|
||
| with tf.device(a.device): | ||
| if ord == "nuc" and isinstance(axis, tuple) and len(axis) != 1: |
There was a problem hiding this comment.
Please, check that keepdims is supported in all cases. I found that this case does not support for keepdims=True.
Please also add a test for this case.
| from nncf.tensor.functions import linalg | ||
|
|
||
|
|
||
| @linalg.norm.register(tf.Tensor) |
There was a problem hiding this comment.
Please add also implementation for ord=0.
| with tf.device(a.device): | ||
| if driver is not None: | ||
| warnings.warn("Driver specifying is not supported in TensorFlow lstsq method") | ||
| if tf.rank(b) == 1: | ||
| b = tf.expand_dims(b, axis=0) | ||
| perm = list(range(tf.rank(b))) | ||
| perm[-1], perm[-2] = perm[-2], perm[-1] | ||
| b = tf.transpose(b, perm=perm) | ||
|
|
||
| return tf.linalg.lstsq(a, b) |
There was a problem hiding this comment.
I run test cases with the code from the suggestion and tests were passed:
Please add test when rank(b) > 2 for example:
import numpy as np
x = [1.0, 2.0, 4.0]
a = np.vstack([x, np.ones_like(x) ** 0]).transpose()
b = np.array([[6.0, 6.0], [8.0, 10.0], [12.0, 18.0]])
np.linalg.lstsq(a, b)[0]
nncf/tensor/functions/tf_numeric.py
Outdated
| with tf.device(a.device): | ||
| if axis is None: | ||
| return tf.reduce_max(a) | ||
| return tf.reduce_max(a, axis=axis, keepdims=keepdim) |
There was a problem hiding this comment.
I would ask you to fix torch in this PR as it is a bug. Thanks
| @numeric.unstack.register(tf.Tensor) | ||
| def _(x: tf.Tensor, axis: int = 0) -> List[tf.Tensor]: | ||
| with tf.device(x.device): | ||
| if not list(x.shape): |
There was a problem hiding this comment.
@AlexanderDokuchaev, could you provide a code snippet when this if will be true?
nncf/tensor/functions/tf_numeric.py
Outdated
| keepdims: bool = False, | ||
| ) -> tf.Tensor: | ||
| numpy_a = np.array(a) | ||
| numpy_median = np.median(numpy_a, axis=axis, keepdims=keepdims) |
There was a problem hiding this comment.
Do you have any evidence that this implementation is more performant than calling np.median?
|
|
||
|
|
||
| @numeric.round.register(tf.Tensor) | ||
| def _(a: tf.Tensor, decimals=0) -> tf.Tensor: |
There was a problem hiding this comment.
| def _(a: tf.Tensor, decimals=0) -> tf.Tensor: | |
| def _(a: tf.Tensor, decimals: int = 0) -> tf.Tensor: |
Changes
tf_numeric.pyandtf_linalg.pyfiles with implementations of methods needed fornncf.Tensorsupport.nncf.Tensor.__ifloordiv__operator fornncf.Tensor.Reason for changes
Currently TensorFlow tensors are not supported by
nncf.Tensor. It prevents #3041 from being done.Related tickets
#3041
Tests
TestTFNNCFTensorOperatorsandTestGPUTFNNCFTensorOperatorsclasses were added totests/tensorflow/test_tensor.py. Some changes were necessary fortests/cross_fw/test_templates/template_test_nncf_tensor.py, mostly related to different device management in TensorFlow.