|
|
|
@ -12,6 +12,7 @@
|
|
|
|
|
# See the License for the specific language governing permissions and
|
|
|
|
|
# limitations under the License.
|
|
|
|
|
|
|
|
|
|
import re
|
|
|
|
|
import numpy as np
|
|
|
|
|
|
|
|
|
|
import mindspore as ms
|
|
|
|
@ -78,9 +79,11 @@ def test_auto_parallel_arithmetic():
|
|
|
|
|
b = Tensor(np.ones([64, 128]), dtype=ms.float32)
|
|
|
|
|
compile_net(net, x, y, b, phase='train')
|
|
|
|
|
strategies = _executor._get_shard_strategy(net)
|
|
|
|
|
expected_strategies = {'Default/network-Net/FloorDiv-op0': [[2, 4], [2, 4]],
|
|
|
|
|
'Default/network-Net/MatMul-op1': [[2, 1], [1, 4]]}
|
|
|
|
|
assert strategies == expected_strategies
|
|
|
|
|
for (k, v) in strategies.items():
|
|
|
|
|
if re.search('FloorDiv-op', k) is not None:
|
|
|
|
|
assert v == [[2, 4], [2, 4]]
|
|
|
|
|
elif re.search('MatMul-op', k) is not None:
|
|
|
|
|
assert v == [[2, 1], [1, 4]]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_auto_parallel_arithmetic_broadcast_both():
|
|
|
|
@ -105,9 +108,11 @@ def test_auto_parallel_arithmetic_broadcast_both():
|
|
|
|
|
b = Tensor(np.ones([1, 64]), dtype=ms.float32)
|
|
|
|
|
compile_net(net, x, y, b, phase='train')
|
|
|
|
|
strategies = _executor._get_shard_strategy(net)
|
|
|
|
|
expected_strategies = {'Default/network-Net/FloorDiv-op0': [[8, 1], [1, 1]],
|
|
|
|
|
'Default/network-Net/MatMul-op1': [[8, 1], [1, 1]]}
|
|
|
|
|
assert strategies == expected_strategies
|
|
|
|
|
for (k, v) in strategies.items():
|
|
|
|
|
if re.search('FloorDiv-op', k) is not None:
|
|
|
|
|
assert v == [[8, 1], [1, 1]]
|
|
|
|
|
elif re.search('MatMul-op', k) is not None:
|
|
|
|
|
assert v == [[8, 1], [1, 1]]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_auto_parallel_arithmetic_broadcast_right():
|
|
|
|
@ -132,9 +137,11 @@ def test_auto_parallel_arithmetic_broadcast_right():
|
|
|
|
|
b = Tensor(np.ones([32]), dtype=ms.float32)
|
|
|
|
|
compile_net(net, x, y, b, phase='train')
|
|
|
|
|
strategies = _executor._get_shard_strategy(net)
|
|
|
|
|
expected_strategies = {'Default/network-Net/FloorDiv-op0': [[4, 2], [2]],
|
|
|
|
|
'Default/network-Net/MatMul-op1': [[4, 1], [1, 2]]}
|
|
|
|
|
assert strategies == expected_strategies
|
|
|
|
|
for (k, v) in strategies.items():
|
|
|
|
|
if re.search('FloorDiv-op', k) is not None:
|
|
|
|
|
assert v == [[4, 2], [2]]
|
|
|
|
|
elif re.search('MatMul-op', k) is not None:
|
|
|
|
|
assert v == [[4, 1], [1, 2]]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_auto_parallel_arithmetic_broadcast_left():
|
|
|
|
@ -159,6 +166,8 @@ def test_auto_parallel_arithmetic_broadcast_left():
|
|
|
|
|
b = Tensor(np.ones([128, 64, 32]), dtype=ms.float32)
|
|
|
|
|
compile_net(net, x, y, b, phase="train")
|
|
|
|
|
strategies = _executor._get_shard_strategy(net)
|
|
|
|
|
expected_strategies = {'Default/network-Net/FloorDiv-op0': [[4, 2], [1, 4, 2]],
|
|
|
|
|
'Default/network-Net/MatMul-op1': [[4, 1], [1, 2]]}
|
|
|
|
|
assert strategies == expected_strategies
|
|
|
|
|
for (k, v) in strategies.items():
|
|
|
|
|
if re.search('FloorDiv-op', k) is not None:
|
|
|
|
|
assert v == [[4, 2], [1, 4, 2]]
|
|
|
|
|
elif re.search('MatMul-op', k) is not None:
|
|
|
|
|
assert v == [[4, 1], [1, 2]]
|
|
|
|
|