Skip to content

Commit

Permalink
Fix PTNNCFCollectorTensorProcessor for cuda + cuda tests for reducers…
Browse files Browse the repository at this point in the history 
…/aggregators
  • Loading branch information
@daniil-lyakhov
daniil-lyakhov committed Sep 28, 2023
1 parent 1b5ff64 commit 60c2804
Show file tree
Hide file tree
Showing 6 changed files with 115 additions and 53 deletions.
21 changes: 13 additions & 8 deletions nncf/common/tensor_statistics/collectors.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ class NNCFCollectorTensorProcessor(ABC):

@staticmethod
@abstractmethod
def reduce_min(x: NNCFTensor, axis: Union[int, tuple, list], keepdims: bool = False) -> NNCFTensor:
def reduce_min(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims: bool = False) -> NNCFTensor:
"""
Computes minimum of elements across dimensions of NNCFTensor.
Expand All @@ -129,7 +129,7 @@ def reduce_min(x: NNCFTensor, axis: Union[int, tuple, list], keepdims: bool = Fa

@staticmethod
@abstractmethod
def reduce_max(x: NNCFTensor, axis: Union[int, tuple, list], keepdims: bool = False) -> NNCFTensor:
def reduce_max(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims: bool = False) -> NNCFTensor:
"""
Computes maximum of elements across dimensions of NNCFTensor.
Expand Down Expand Up @@ -174,7 +174,7 @@ def max(x1: NNCFTensor, x2: NNCFTensor) -> NNCFTensor:

@staticmethod
@abstractmethod
def mean(x: NNCFTensor, axis: Union[int, tuple, list], keepdims=False) -> NNCFTensor:
def mean(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims=False) -> NNCFTensor:
"""
Computes the mean of elements across given dimensions of NNCFTensor.
Expand All @@ -187,7 +187,7 @@ def mean(x: NNCFTensor, axis: Union[int, tuple, list], keepdims=False) -> NNCFTe

@staticmethod
@abstractmethod
def median(x: NNCFTensor, axis: Union[int, tuple, list], keepdims=False) -> NNCFTensor:
def median(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims=False) -> NNCFTensor:
"""
Computes the median of elements across given dimensions of NNCFTensor.
Expand All @@ -200,7 +200,9 @@ def median(x: NNCFTensor, axis: Union[int, tuple, list], keepdims=False) -> NNCF

@classmethod
@abstractmethod
def masked_mean(cls, x: NNCFTensor, axis: Union[int, tuple, list], mask: NNCFTensor, keepdims=False) -> NNCFTensor:
def masked_mean(
cls, x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], mask: NNCFTensor, keepdims=False
) -> NNCFTensor:
"""
Computes the masked mean of elements across given dimensions of NNCFTensor.
Expand All @@ -216,7 +218,7 @@ def masked_mean(cls, x: NNCFTensor, axis: Union[int, tuple, list], mask: NNCFTen
@classmethod
@abstractmethod
def masked_median(
cls, x: NNCFTensor, axis: Union[int, tuple, list], mask: NNCFTensor, keepdims=False
cls, x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], mask: NNCFTensor, keepdims=False
) -> NNCFTensor:
"""
Computes the masked median of elements across given dimensions of NNCFTensor.
Expand Down Expand Up @@ -275,7 +277,10 @@ def sum(tensor: NNCFTensor) -> TensorElementsType:
@staticmethod
@abstractmethod
def quantile(
tensor: NNCFTensor, quantile: Union[float, List[float]], axis: Union[int, tuple, list], keepdims: bool = False
tensor: NNCFTensor,
quantile: Union[float, List[float]],
axis: Union[int, Tuple[int, ...], List[int]],
keepdims: bool = False,
) -> List[TensorElementsType]:
"""
Compute the quantile(s) of the data along the specified axis.
Expand All @@ -295,7 +300,7 @@ def percentile(
cls,
tensor: NNCFTensor,
percentile: Union[float, List[float]],
axis: Union[int, tuple, list],
axis: Union[int, Tuple[int, ...], List[int]],
keepdims: bool = False,
) -> List[TensorElementsType]:
"""
Expand Down
27 changes: 19 additions & 8 deletions nncf/onnx/statistics/collectors.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,11 +33,11 @@ class ONNXNNCFCollectorTensorProcessor(NNCFCollectorTensorProcessor):
"""

@staticmethod
def reduce_min(x: NNCFTensor, axis: Union[int, Tuple, list], keepdims: bool = False) -> NNCFTensor:
def reduce_min(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims: bool = False) -> NNCFTensor:
return ONNXNNCFTensor(np.amin(x.tensor, axis=axis, keepdims=keepdims))

@staticmethod
def reduce_max(x: NNCFTensor, axis: Union[int, Tuple, list], keepdims: bool = False) -> NNCFTensor:
def reduce_max(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims: bool = False) -> NNCFTensor:
return ONNXNNCFTensor(np.amax(x.tensor, axis=axis, keepdims=keepdims))

@staticmethod
Expand All @@ -53,16 +53,20 @@ def max(x1: NNCFTensor, x2: NNCFTensor) -> NNCFTensor:
return ONNXNNCFTensor(np.maximum(x1.tensor, x2.tensor))

@staticmethod
def mean(x: NNCFTensor, axis: Union[int, Tuple, list], keepdims=False) -> NNCFTensor:
def mean(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims=False) -> NNCFTensor:
return ONNXNNCFTensor(np.mean(x.tensor, axis=axis, keepdims=keepdims))

@staticmethod
def median(x: NNCFTensor, axis: Union[int, Tuple, list], keepdims=False) -> NNCFTensor:
def median(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims=False) -> NNCFTensor:
return ONNXNNCFTensor(np.median(x.tensor, axis=axis, keepdims=keepdims))

@classmethod
def masked_mean(
cls, x: NNCFTensor, axis: Optional[Union[int, Tuple, list]], mask: Optional[NNCFTensor], keepdims: bool = False
cls,
x: NNCFTensor,
axis: Optional[Union[int, Tuple[int, ...], List[int]]],
mask: Optional[NNCFTensor],
keepdims: bool = False,
) -> NNCFTensor:
if mask is None:
return cls.mean(x, axis=axis, keepdims=keepdims)
Expand All @@ -71,7 +75,11 @@ def masked_mean(

@classmethod
def masked_median(
cls, x: NNCFTensor, axis: Optional[Union[int, Tuple, list]], mask: Optional[NNCFTensor], keepdims: bool = False
cls,
x: NNCFTensor,
axis: Optional[Union[int, Tuple[int, ...], List[int]]],
mask: Optional[NNCFTensor],
keepdims: bool = False,
) -> NNCFTensor:
if mask is None:
return cls.median(x, axis=axis, keepdims=keepdims)
Expand Down Expand Up @@ -105,7 +113,10 @@ def sum(tensor: NNCFTensor) -> TensorElementsType:

@staticmethod
def quantile(
tensor: NNCFTensor, quantile: Union[float, List[float]], axis: Union[int, Tuple, list], keepdims: bool = False
tensor: NNCFTensor,
quantile: Union[float, List[float]],
axis: Union[int, Tuple[int, ...], List[int]],
keepdims: bool = False,
) -> List[TensorElementsType]:
result = np.quantile(tensor.tensor, quantile, axis, keepdims=keepdims)
return [ONNXNNCFTensor(x) for x in result]
Expand All @@ -115,7 +126,7 @@ def percentile(
cls,
tensor: NNCFTensor,
percentile: Union[float, List[float]],
axis: Union[int, Tuple, list],
axis: Union[int, Tuple[int, ...], List[int]],
keepdims: bool = False,
) -> List[TensorElementsType]:
raise NotImplementedError()
Expand Down
27 changes: 19 additions & 8 deletions nncf/openvino/statistics/collectors.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,11 +50,11 @@ class OVNNCFCollectorTensorProcessor(NNCFCollectorTensorProcessor):
"""

@staticmethod
def reduce_min(x: NNCFTensor, axis: Union[int, Tuple], keepdims: bool = True) -> NNCFTensor:
def reduce_min(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims: bool = True) -> NNCFTensor:
return OVNNCFTensor(np.amin(x.tensor, axis=axis, keepdims=keepdims))

@staticmethod
def reduce_max(x: NNCFTensor, axis: Union[int, Tuple], keepdims: bool = True) -> NNCFTensor:
def reduce_max(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims: bool = True) -> NNCFTensor:
return OVNNCFTensor(np.amax(x.tensor, axis=axis, keepdims=keepdims))

@staticmethod
Expand All @@ -70,16 +70,20 @@ def max(x1: NNCFTensor, x2: NNCFTensor) -> NNCFTensor:
return OVNNCFTensor(np.maximum(x1.tensor, x2.tensor))

@staticmethod
def mean(x: NNCFTensor, axis: Union[int, Tuple], keepdims: bool = False) -> NNCFTensor:
def mean(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims: bool = False) -> NNCFTensor:
return OVNNCFTensor(np.mean(x.tensor, axis=axis, keepdims=keepdims))

@staticmethod
def median(x: NNCFTensor, axis: Union[int, Tuple, list], keepdims: bool = False) -> NNCFTensor:
def median(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims: bool = False) -> NNCFTensor:
return OVNNCFTensor(np.median(x.tensor, axis=axis, keepdims=keepdims))

@classmethod
def masked_mean(
cls, x: NNCFTensor, axis: Optional[Union[int, Tuple, list]], mask: Optional[NNCFTensor], keepdims: bool = False
cls,
x: NNCFTensor,
axis: Optional[Union[int, Tuple[int, ...], List[int]]],
mask: Optional[NNCFTensor],
keepdims: bool = False,
) -> NNCFTensor:
if mask is None:
return cls.mean(x, axis=axis, keepdims=keepdims)
Expand All @@ -91,7 +95,11 @@ def masked_mean(

@classmethod
def masked_median(
cls, x: NNCFTensor, axis: Optional[Union[int, Tuple, list]], mask: Optional[NNCFTensor], keepdims: bool = False
cls,
x: NNCFTensor,
axis: Optional[Union[int, Tuple[int, ...], List[int]]],
mask: Optional[NNCFTensor],
keepdims: bool = False,
) -> NNCFTensor:
if mask is None:
return cls.median(x, axis=axis, keepdims=keepdims)
Expand Down Expand Up @@ -140,7 +148,10 @@ def sum(tensor: NNCFTensor) -> TensorElementsType:

@staticmethod
def quantile(
tensor: NNCFTensor, quantile: Union[float, List[float]], axis: Union[int, Tuple, list], keepdims: bool = False
tensor: NNCFTensor,
quantile: Union[float, List[float]],
axis: Union[int, Tuple[int, ...], List[int]],
keepdims: bool = False,
) -> List[NNCFTensor]:
result = np.quantile(tensor.tensor, quantile, axis, keepdims=keepdims)
return [OVNNCFTensor(x) for x in result]
Expand All @@ -150,7 +161,7 @@ def percentile(
cls,
tensor: NNCFTensor,
percentile: Union[float, List[float]],
axis: Union[int, tuple, list],
axis: Union[int, Tuple[int, ...], List[int]],
keepdims: bool = False,
) -> List[TensorElementsType]:
quantile = np.true_divide(percentile, 100)
Expand Down
21 changes: 13 additions & 8 deletions nncf/tensorflow/tensor_statistics/collectors.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,11 @@ class TFNNCFCollectorTensorProcessor(NNCFCollectorTensorProcessor):
"""

@staticmethod
def reduce_min(x: NNCFTensor, axis: Union[int, tuple, list], keepdims: bool = False) -> NNCFTensor:
def reduce_min(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims: bool = False) -> NNCFTensor:
return TFNNCFTensor(tf.reduce_min(x.tensor, axis=axis, keepdims=keepdims))

@staticmethod
def reduce_max(x: NNCFTensor, axis: Union[int, tuple, list], keepdims: bool = False) -> NNCFTensor:
def reduce_max(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims: bool = False) -> NNCFTensor:
return TFNNCFTensor(tf.reduce_max(x.tensor, axis=axis, keepdims=keepdims))

@staticmethod
Expand All @@ -57,20 +57,22 @@ def max(x1: tf.Tensor, x2: tf.Tensor) -> NNCFTensor:
return TFNNCFTensor(tf.math.maximum(x1.tensor, x2.tensor))

@staticmethod
def mean(x: NNCFTensor, axis: Union[int, tuple, list], keepdims=False) -> NNCFTensor:
def mean(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims=False) -> NNCFTensor:
return TFNNCFTensor(tf.math.reduce_mean(x.tensor, axis=axis, keepdims=keepdims))

@staticmethod
def median(x: NNCFTensor, axis: Union[int, tuple, list], keepdims=False) -> NNCFTensor:
def median(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims=False) -> NNCFTensor:
raise NotImplementedError()

@classmethod
def masked_mean(cls, x: NNCFTensor, axis: Union[int, tuple, list], mask: NNCFTensor, keepdims=False) -> NNCFTensor:
def masked_mean(
cls, x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], mask: NNCFTensor, keepdims=False
) -> NNCFTensor:
raise NotImplementedError()

@classmethod
def masked_median(
cls, x: NNCFTensor, axis: Union[int, tuple, list], mask: NNCFTensor, keepdims=False
cls, x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], mask: NNCFTensor, keepdims=False
) -> NNCFTensor:
raise NotImplementedError()

Expand Down Expand Up @@ -105,7 +107,10 @@ def sum(tensor: NNCFTensor) -> TensorElementsType:

@staticmethod
def quantile(
tensor: NNCFTensor, quantile: Union[float, List[float]], axis: Union[int, tuple, list], keepdims: bool = False
tensor: NNCFTensor,
quantile: Union[float, List[float]],
axis: Union[int, Tuple[int, ...], List[int]],
keepdims: bool = False,
) -> List[NNCFTensor]:
raise NotImplementedError()

Expand All @@ -114,7 +119,7 @@ def percentile(
cls,
tensor: NNCFTensor,
percentile: Union[float, List[float]],
axis: Union[int, tuple, list],
axis: Union[int, Tuple[int, ...], List[int]],
keepdims: bool = False,
) -> List[TensorElementsType]:
raise NotImplementedError()
Expand Down
42 changes: 26 additions & 16 deletions nncf/torch/tensor_statistics/collectors.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,11 +50,11 @@ class PTNNCFCollectorTensorProcessor(NNCFCollectorTensorProcessor):
"""

@staticmethod
def reduce_min(x: NNCFTensor, axis: Union[int, tuple, list], keepdims: bool = False) -> NNCFTensor:
def reduce_min(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims: bool = False) -> NNCFTensor:
return PTNNCFTensor(torch.amin(x.tensor, dim=axis, keepdim=keepdims))

@staticmethod
def reduce_max(x: NNCFTensor, axis: Union[int, tuple, list], keepdims: bool = False) -> NNCFTensor:
def reduce_max(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims: bool = False) -> NNCFTensor:
return PTNNCFTensor(torch.amax(x.tensor, dim=axis, keepdim=keepdims))

@staticmethod
Expand All @@ -72,38 +72,44 @@ def max(cls, *args) -> NNCFTensor:
return cls.reduce_max(stacked, axis=0, keepdims=False)

@staticmethod
def mean(x: NNCFTensor, axis: Union[int, tuple, list], keepdims=False) -> NNCFTensor:
def mean(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims=False) -> NNCFTensor:
return PTNNCFTensor(x.tensor.mean(dim=axis, keepdim=keepdims))

@staticmethod
def median(x: NNCFTensor, axis: Union[int, tuple, list], keepdims=False) -> NNCFTensor:
def median(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], keepdims=False) -> NNCFTensor:
# See https://github.com/pytorch/pytorch/issues/61582
if not isinstance(axis, int):
return PTNNCFTensor(torch.tensor(np.median(x.tensor.detach().cpu().numpy(), axis=axis, keepdims=keepdims)))
device = x.tensor.device
result = torch.tensor(np.median(x.tensor.detach().cpu().numpy(), axis=axis, keepdims=keepdims))
return PTNNCFTensor(result.type(x.tensor.dtype).to(device))
return PTNNCFTensor(torch.quantile(x.tensor, q=0.5, dim=axis, keepdim=keepdims).values)

@classmethod
def masked_mean(cls, x: NNCFTensor, axis: Union[int, tuple], mask: NNCFTensor, keepdims=False) -> NNCFTensor:
def masked_mean(
cls, x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], mask: NNCFTensor, keepdims=False
) -> NNCFTensor:
if mask is None:
return cls.mean(x, axis=axis, keepdims=keepdims)
masked_x = np.ma.array(x.tensor.detach().cpu().numpy(), mask=mask.tensor)
device = x.tensor.device
masked_x = np.ma.array(x.tensor.detach().cpu().numpy(), mask=mask.tensor.detach().cpu().numpy())
result = np.ma.mean(masked_x, axis=axis, keepdims=keepdims).astype(masked_x.dtype)
if isinstance(result, np.ma.MaskedArray):
return PTNNCFTensor(torch.tensor(result.data))
return PTNNCFTensor(torch.tensor(result))
result = result.data
return PTNNCFTensor(torch.tensor(result).to(device=device))

@classmethod
def masked_median(
cls, x: NNCFTensor, axis: Union[int, tuple, list], mask: NNCFTensor, keepdims=False
cls, x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], mask: NNCFTensor, keepdims=False
) -> NNCFTensor:
# Implemented in numy as torch.masked.median is not implemented yet
if mask is None:
return cls.median(x, axis=axis, keepdims=keepdims)
device = x.tensor.device
masked_x = np.ma.array(x.tensor.detach().cpu().numpy(), mask=mask.tensor.detach().cpu().numpy())
result = np.ma.median(masked_x, axis=axis, keepdims=keepdims).astype(masked_x.dtype)
if isinstance(result, np.ma.MaskedArray):
return PTNNCFTensor(torch.tensor(result.data))
return PTNNCFTensor(torch.tensor(result))
result = result.data
return PTNNCFTensor(torch.tensor(result).to(device=device))

@staticmethod
def mean_per_channel(x: NNCFTensor, axis: int) -> NNCFTensor:
Expand Down Expand Up @@ -148,24 +154,28 @@ def sum(tensor: NNCFTensor) -> TensorElementsType:

@staticmethod
def quantile(
tensor: NNCFTensor, quantile: Union[float, List[float]], axis: Union[int, tuple, list], keepdims: bool = False
tensor: NNCFTensor,
quantile: Union[float, List[float], np.ndarray],
axis: Union[int, Tuple[int, ...], List[int]],
keepdims: bool = False,
) -> List[NNCFTensor]:
device = tensor.device
# See https://github.com/pytorch/pytorch/issues/61582
if not isinstance(axis, int):
result = torch.tensor(
np.quantile(tensor.tensor.detach().cpu().numpy(), q=quantile, axis=axis, keepdims=keepdims)
)
else:
result = torch.quantile(tensor.tensor, torch.tensor(quantile).type(tensor.tensor.dtype), axis, keepdims)
result = result.type(tensor.tensor.dtype)
result = result.type(tensor.tensor.dtype).to(device)
return [PTNNCFTensor(x) for x in result]

@classmethod
def percentile(
cls,
tensor: NNCFTensor,
percentile: Union[float, List[float]],
axis: Union[int, tuple, list],
percentile: Union[float, List[float], np.ndarray],
axis: Union[int, Tuple[int, ...], List[int]],
keepdims: bool = False,
) -> List[TensorElementsType]:
quantile = np.true_divide(percentile, 100)
Expand Down
Loading

0 comments on commit 60c2804

Please sign in to comment.