forked from mindspore-Ecosystem/mindspore
28 lines
1.0 KiB
Python
28 lines
1.0 KiB
Python
import numpy as np
|
|
from mindspore import Tensor
|
|
from mindspore.parallel.nn.layers import FixedSparseAttention
|
|
import mindspore.context as context
|
|
|
|
context.set_context(device_target="Ascend")
|
|
|
|
|
|
def test_net():
|
|
np.random.seed(0)
|
|
bs = 2 # batch size
|
|
heads = 2
|
|
seq_len = 1024 # this op is designed for seq_len = 1024
|
|
size_per_head = 128 # maximum size per head value is 128
|
|
|
|
block_size = 64 # block size is designed to be 64
|
|
fixed_sparse = FixedSparseAttention(bs, heads, size_per_head, block_size)
|
|
q = np.random.rand(bs, seq_len, heads * size_per_head)
|
|
q = q.astype(np.float16)
|
|
k = np.random.rand(bs, seq_len, heads * size_per_head)
|
|
k = k.astype(np.float16)
|
|
v = np.random.rand(bs, seq_len, heads * size_per_head)
|
|
v = v.astype(np.float16)
|
|
attention_mask = np.ones((bs, seq_len, seq_len), dtype=np.float32)
|
|
out = fixed_sparse(Tensor(q), Tensor(k), Tensor(v), Tensor(attention_mask))
|
|
out_np = out.asnumpy()
|
|
print("local output: ", out_np[0, 0])
|