forked from huawei/mindspore2022
111 lines
4.0 KiB
Python
111 lines
4.0 KiB
Python
import pytest
|
|
import numpy as np
|
|
|
|
import mindspore.context as context
|
|
import mindspore.nn as nn
|
|
from mindspore import Tensor
|
|
from mindspore.common.api import ms_function
|
|
from mindspore.ops import operations as P
|
|
|
|
|
|
class Net(nn.Cell):
|
|
def __init__(self, seq_dim, batch_dim):
|
|
super(Net, self).__init__()
|
|
self.reverse_sequence = P.ReverseSequence(
|
|
seq_dim=seq_dim, batch_dim=batch_dim)
|
|
|
|
@ms_function
|
|
def construct(self, x, seq_lengths):
|
|
return self.reverse_sequence(x, seq_lengths)
|
|
|
|
|
|
@pytest.mark.level0
|
|
@pytest.mark.platform_x86_gpu_training
|
|
@pytest.mark.env_onecard
|
|
def test_net_int8():
|
|
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
|
|
x = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]).astype(np.int8)
|
|
seq_lengths = np.array([1, 2, 3]).astype(np.int32)
|
|
seq_dim = 0
|
|
batch_dim = 1
|
|
net = Net(seq_dim, batch_dim)
|
|
output = net(Tensor(x), Tensor(seq_lengths))
|
|
expected = np.array([[1, 5, 9], [4, 2, 6], [7, 8, 3]]).astype(np.int8)
|
|
assert np.array_equal(output.asnumpy(), expected)
|
|
|
|
|
|
@pytest.mark.level0
|
|
@pytest.mark.platform_x86_gpu_training
|
|
@pytest.mark.env_onecard
|
|
def test_net_int32():
|
|
context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")
|
|
x = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]).astype(np.int32)
|
|
seq_lengths = np.array([1, 2, 3]).astype(np.int64)
|
|
seq_dim = 1
|
|
batch_dim = 0
|
|
net = Net(seq_dim, batch_dim)
|
|
output = net(Tensor(x), Tensor(seq_lengths))
|
|
expected = np.array([[1, 2, 3], [5, 4, 6], [9, 8, 7]]).astype(np.int32)
|
|
assert np.array_equal(output.asnumpy(), expected)
|
|
|
|
|
|
@pytest.mark.level0
|
|
@pytest.mark.platform_x86_gpu_training
|
|
@pytest.mark.env_onecard
|
|
def test_net_float32():
|
|
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
|
|
x = np.array([[[[1, 2], [3, 4]],
|
|
[[5, 6], [7, 8]],
|
|
[[9, 10], [11, 12]],
|
|
[[13, 14], [15, 16]]],
|
|
[[[17, 18], [19, 20]],
|
|
[[21, 22], [23, 24]],
|
|
[[25, 26], [27, 28]],
|
|
[[29, 30], [31, 21]]]]).astype(np.float32)
|
|
seq_lengths = np.array([2, 2, 2, 2]).astype(np.int64)
|
|
seq_dim = 0
|
|
batch_dim = 1
|
|
net = Net(seq_dim, batch_dim)
|
|
output = net(Tensor(x), Tensor(seq_lengths))
|
|
expected = np.array([[[[17., 18.], [19., 20.]],
|
|
[[21., 22.], [23., 24.]],
|
|
[[25., 26.], [27., 28.]],
|
|
[[29., 30.], [31., 21.]]],
|
|
[[[1., 2.], [3., 4.]],
|
|
[[5., 6.], [7., 8.]],
|
|
[[9., 10.], [11., 12.]],
|
|
[[13., 14.], [15., 16.]]]]).astype(np.float32)
|
|
assert np.array_equal(output.asnumpy(), expected)
|
|
|
|
|
|
@pytest.mark.level0
|
|
@pytest.mark.platform_x86_gpu_training
|
|
@pytest.mark.env_onecard
|
|
def test_net_float64_0_dim():
|
|
"""
|
|
Test added to test for 0 seq len edge case
|
|
"""
|
|
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
|
|
x = np.array([[[[1, 2], [3, 4]],
|
|
[[5, 6], [7, 8]],
|
|
[[9, 10], [11, 12]],
|
|
[[13, 14], [15, 16]]],
|
|
[[[17, 18], [19, 20]],
|
|
[[21, 22], [23, 24]],
|
|
[[25, 26], [27, 28]],
|
|
[[29, 30], [31, 21]]]]).astype(np.float32)
|
|
seq_lengths = np.array([2, 2, 0, 0]).astype(np.int64)
|
|
seq_dim = 2
|
|
batch_dim = 1
|
|
net = Net(seq_dim, batch_dim)
|
|
output = net(Tensor(x), Tensor(seq_lengths))
|
|
expected = np.array([[[[3., 4.], [1., 2.]],
|
|
[[7., 8.], [5., 6.]],
|
|
[[9., 10.], [11., 12.]],
|
|
[[13., 14.], [15., 16.]]],
|
|
[[[19., 20.], [17., 18.]],
|
|
[[23., 24.], [21., 22.]],
|
|
[[25., 26.], [27., 28.]],
|
|
[[29., 30.], [31., 21.]]]]).astype(np.float32)
|
|
assert np.array_equal(output.asnumpy(), expected)
|