def test_optimizer_clone_weight2():
class Net(nn.Cell):
def __init__(self, strategy1, strategy2, weight):
super().__init__()
self.weight = Parameter(weight, "w1")
self.matmul = P.MatMul(transpose_a=False,
transpose_b=True).set_strategy(strategy1)
self.relu = P.ReLU().set_strategy(strategy2)
def construct(self, x):
out = self.matmul(x, self.weight)
out = self.relu(out)
return out
context.set_auto_parallel_context(device_num=4, global_rank=0)
strategy1 = ((2, 1), (2, 1))
strategy2 = ((4, 1), )
strategy3 = ((4, 1), (4, 1))
x = Tensor(np.ones([64, 32]), dtype=ms.float32)
weight = Tensor(np.ones([64, 32]), dtype=ms.float32)
b = Tensor(np.ones([64, 64]), dtype=ms.float32)
net = Net(strategy1, strategy2, weight)
optimizer = AdamWeightDecay(net.trainable_params())
net_with_loss = NetWithLoss(net, strategy3)
train_net = TrainOneStepCell(net_with_loss, optimizer)
context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")
_Executor().compile(train_net, x, b)
def compile(net, x, b):
net.set_auto_parallel()
_Executor().compile(net, x, b)