natas_hw
/
openvino
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[PT FE] ScaledDotProductAttention fix issue with other float dtypes (#21638) 

* Fix scale input

* Improve torch scaled dot tests
This commit is contained in:
Mateusz Mikolajczyk 2023-12-18 11:32:54 +01:00 committed by GitHub
parent e224e4f4f0
commit be476519fb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 11 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/frontends/pytorch/src/op/scaled_dot_product_attention.cpp
View File

@ -32,7 +32,7 @@ std::shared_ptr<ov::Node> translate_scaled_dot_product_attention_common(const No
inputs.push_back(context.mark_node(std::make_shared<v1::ConvertLike>(zero, query))); inputs.push_back(context.mark_node(std::make_shared<v1::ConvertLike>(zero, query)));
} }
if (!context.input_is_none(6)) if (!context.input_is_none(6))
inputs.push_back(context.get_input(6)); inputs.push_back(context.mark_node(std::make_shared<v1::ConvertLike>(context.get_input(6), query)));
return context.mark_node(std::make_shared<v13::ScaledDotProductAttention>(inputs, is_causal)); return context.mark_node(std::make_shared<v13::ScaledDotProductAttention>(inputs, is_causal));
} }

19
tests/layer_tests/pytorch_tests/test_scaled_dot_product_attention.py
View File

@ -9,20 +9,20 @@ from pytorch_layer_test_class import PytorchLayerTest
class TestScaledDotProductAttention(PytorchLayerTest): class TestScaledDotProductAttention(PytorchLayerTest):
def _prepare_input(self): def _prepare_input(self, dtype):
return (np.random.randn(1, 2, 8, 4).astype(np.float32), np.random.randn(1, 2, 8, 4).astype(np.float32), np.random.randn(1, 2, 8, 4).astype(np.float32)) return (np.random.randn(1, 2, 8, 4).astype(dtype), np.random.randn(1, 2, 8, 4).astype(dtype), np.random.randn(1, 2, 8, 4).astype(dtype))
def create_model(self, mask, is_causal, scale): def create_model(self, mask, is_causal, scale, dtype):
import torch.nn.functional as F import torch.nn.functional as F
import torch import torch
class aten_scaled_dot_product_atten(torch.nn.Module): class aten_scaled_dot_product_atten(torch.nn.Module):
def __init__(self, mask=False, is_causal=False, scale=False) -> None: def __init__(self, mask=False, is_causal=False, scale=False, dtype=np.float32) -> None:
super().__init__() super().__init__()
self.mask = None if not mask else torch.from_numpy( self.mask = None if not mask else torch.from_numpy(
np.random.randint(0, 2, (8, 8)).astype(np.float32)) np.random.randint(0, 2, (8, 8)).astype(dtype))
self.is_causal = is_causal self.is_causal = is_causal
if is_causal and mask: if is_causal and mask:
self.mask.to(torch.bool) self.mask.to(torch.bool)
@ -36,13 +36,14 @@ class TestScaledDotProductAttention(PytorchLayerTest):
ref_net = None ref_net = None
return aten_scaled_dot_product_atten(mask, is_causal, scale), ref_net, 'aten::scaled_dot_product_attention' return aten_scaled_dot_product_atten(mask, is_causal, scale, dtype), ref_net, 'aten::scaled_dot_product_attention'
@pytest.mark.nightly @pytest.mark.nightly
@pytest.mark.precommit @pytest.mark.precommit
@pytest.mark.precommit_fx_backend @pytest.mark.precommit_fx_backend
@pytest.mark.parametrize(['mask', 'is_causal'], [(False, False), (False, True), (True, True), (True, False)]) @pytest.mark.parametrize(['mask', 'is_causal'], [(False, False), (False, True), (True, True), (True, False)])
@pytest.mark.parametrize('scale', [False, True]) @pytest.mark.parametrize('scale', [False, True])
def test_scaled_dot_product_atten(self, ie_device, precision, ir_version, mask, is_causal, scale): @pytest.mark.parametrize("dtype", (np.float32, np.float64))
self._test(*self.create_model(mask, is_causal, scale), def test_scaled_dot_product_atten(self, ie_device, precision, ir_version, mask, is_causal, scale, dtype):
ie_device, precision, ir_version) self._test(*self.create_model(mask, is_causal, scale, dtype),
ie_device, precision, ir_version, kwargs_to_prepare_input={"dtype": dtype})