Skip to content

Commit

Permalink
[Torch] Workaroud for torch.quantile function (openvinotoolkit#2240)
Browse files Browse the repository at this point in the history
### Changes

torch quantile is always using numpy realization

### Reason for changes

Bug pytorch/pytorch#64947

third party test 237 has failed

```
       RuntimeError: quantile() input tensor is too large
 ```
### Related tickets


### Tests
Build test third party 239 is finished successfully
  • Loading branch information
daniil-lyakhov authored Nov 2, 2023
1 parent f8bdbc7 commit 2944eb4
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 6 deletions.
13 changes: 7 additions & 6 deletions nncf/torch/tensor_statistics/collectors.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ def median(x: NNCFTensor, axis: Union[int, Tuple[int, ...], List[int]], 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))
return PTNNCFCollectorTensorProcessor.quantile(x, quantile=[0.5], axis=axis, keepdims=keepdims)[0]

@classmethod
def masked_mean(
Expand Down Expand Up @@ -173,17 +173,18 @@ def quantile(
) -> 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:
# https://github.com/pytorch/pytorch/issues/64947
if len(tensor.tensor) <= 16_000_000 and isinstance(axis, int):
result = torch.quantile(
tensor.tensor,
torch.tensor(quantile, dtype=tensor.tensor.dtype, device=tensor.tensor.device),
axis,
keepdims,
)
else:
result = torch.tensor(
np.quantile(tensor.tensor.detach().cpu().numpy(), q=quantile, axis=axis, keepdims=keepdims)
)
result = result.type(tensor.tensor.dtype).to(device)
return [PTNNCFTensor(x) for x in result]

Expand Down
21 changes: 21 additions & 0 deletions tests/torch/ptq/test_reducers_and_aggregators.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,3 +101,24 @@ def get_nncf_tensor(self, x: np.array, dtype: Optional[Dtype] = None):
def all_close(self, val: torch.Tensor, ref) -> bool:
assert val.is_cuda
return super().all_close(val, ref)


@pytest.mark.parametrize("device", ["cuda", "cpu"])
@pytest.mark.parametrize("size,ref", [(16_000_000, 1_600_000.8750), (17_000_000, 1_700_000.7500)])
def test_quantile_percentile_function(device, size, ref):
tensor = PTNNCFTensor(torch.arange(1, size, 1).float().to(device))
res_quantile = PTNNCFCollectorTensorProcessor.quantile(tensor, [0.1], axis=0)
res_percentile = PTNNCFCollectorTensorProcessor.percentile(tensor, [10], axis=0)
assert len(res_quantile) == len(res_percentile) == 1
for tensor in [res_quantile[0].tensor, res_percentile[0].tensor]:
assert tensor == ref
assert tensor.is_cuda == (device == "cuda")


@pytest.mark.parametrize("device", ["cuda", "cpu"])
@pytest.mark.parametrize("size,ref", [(16_000_000, 8_000_000), (17_000_000, 8_500_000)])
def test_median_function(device, size, ref):
tensor = PTNNCFTensor(torch.arange(1, size, 1).float().to(device))
res = PTNNCFCollectorTensorProcessor.median(tensor, axis=0)
assert res.tensor == ref
assert res.tensor.is_cuda == (device == "cuda")

0 comments on commit 2944eb4

Please sign in to comment.