mindspore/tests/st/auto_parallel/onehot_model_parallel.py

155 lines
5.4 KiB
Python

# Copyright 2019 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
import os
import numpy as np
import mindspore as ms
import mindspore.communication.management as distributedTool
import mindspore.context as context
from mindspore.common.tensor import Tensor
from mindspore.nn import Cell
from mindspore.ops import operations as P
device_num = 2
device_id = int(os.getenv('DEVICE_ID'))
rank_id = 0
def setup_module():
global device_num
global rank_id
np.random.seed(0)
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
context.set_context(device_id=device_id)
distributedTool.init()
device_num = distributedTool.get_group_size()
rank_id = distributedTool.get_rank()
context.set_auto_parallel_context(device_num=device_num,
global_rank=rank_id)
def teardown_module():
distributedTool.release()
class Onehot(Cell):
def __init__(self, axis=-1, depth=1, on_value=1.0, off_value=0.0, strategy=None):
super(Onehot, self).__init__()
trans_stra = None
if strategy:
trans_stra = (strategy[0],)
self.onehot = P.OneHot().shard(strategy=strategy)
self.depth = depth
self.on_value = Tensor(on_value, ms.float32)
self.off_value = Tensor(off_value, ms.float32)
self.transpose = P.Transpose().shard(strategy=trans_stra)
self.sub = P.Sub().shard(strategy=((1, 1), (1, 1)))
self.axis = axis
def construct(self, input_, indices):
x = self.onehot(indices, self.depth, self.on_value, self.off_value)
x = self.transpose(x, (1, 0))
x = self.sub(input_, x)
return x
class DataGenerator():
def get_parallel_blocks(self, input_, strategy):
blocks = [input_]
i = 0
for stra in strategy:
temp = []
while blocks:
block = blocks.pop(0)
temp.extend(np.split(block, stra, axis=i))
blocks.extend(temp)
i += 1
return blocks
def generate_data(self, shape):
data = np.random.rand(*shape)
return data
def input_data(self, shape):
data = (self.generate_data(shape) * 2).astype(np.float32)
stra = [1] * len(shape)
stra[0] = device_num
datas = self.get_parallel_blocks(data, stra)
return Tensor(data), Tensor(datas[rank_id])
def label_data(self, shape, classes):
data = (self.generate_data(shape) * (classes - 1)).astype(np.int32)
stra = [1] * len(shape)
stra[0] = device_num
datas = self.get_parallel_blocks(data, stra)
return Tensor(data), Tensor(datas[rank_id])
class OneHotFactory:
def __init__(self, batch_size, classes, on_value=1.0, off_value=0.0, axis=None, strategy=None):
data_gen = DataGenerator()
self.input_full, self.input_part = data_gen.input_data((classes, batch_size))
self.label_full, self.label_part = data_gen.label_data((batch_size,), classes)
self.depth = classes
self.on_value = on_value
self.off_value = off_value
self.axis = axis
self.strategy = strategy
def forward_mindspore_single_impl(self):
net = Onehot(axis=self.axis,
depth=self.depth,
on_value=self.on_value,
off_value=self.off_value)
out = net(self.input_full, self.label_full)
return out
def forward_mindspore_parallel_impl(self):
context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
net = Onehot(axis=self.axis,
depth=self.depth,
on_value=self.on_value,
off_value=self.off_value, strategy=self.strategy)
out = net.compile_and_run(self.input_full, self.label_full)
return out
def forward_cmp(self):
out_mindspore_single = self.forward_mindspore_single_impl().asnumpy()
context.reset_auto_parallel_context()
out_mindspore_parallel = self.forward_mindspore_parallel_impl().asnumpy()
context.reset_auto_parallel_context()
assert np.allclose(out_mindspore_single, out_mindspore_parallel, 0.0001, 0.0001)
def test_reid_onehot_forward_int32_128_depth1024_model_parallel():
fact = OneHotFactory(batch_size=128,
classes=1024,
on_value=1.000000,
off_value=0.000000,
axis=-1,
strategy=((1, device_num), (), ()))
fact.forward_cmp()
def test_reid_onehot_forward_int32_1024_depth128_model_parallel():
fact = OneHotFactory(batch_size=1024,
classes=128,
on_value=1.000000,
off_value=0.000000,
axis=-1,
strategy=((1, device_num), (), ()))
fact.forward_cmp()