def test_parser_switch_layer_switch_in_bprop():
class OneInputBprop(nn.Cell):
def __init__(self, funcs):
super(OneInputBprop, self).__init__()
self.op = P.ReLU()
self.funcs = funcs
def construct(self, i, x):
return self.op(x)
def bprop(self, i, x, out, dout):
return i, self.funcs[i](x, dout)
class Add(nn.Cell):
def __init__(self):
super().__init__()
self.op = P.TensorAdd()
def construct(self, x, y):
return self.op(x, y)
class Mul(nn.Cell):
def __init__(self):
super().__init__()
self.op = P.Mul()
def construct(self, x, y):
return self.op(x, y)
func1 = Add()
func2 = Mul()
funcs = (func1, func2)
net = OneInputBprop(funcs)
input1 = Tensor(np.ones([2, 2]).astype(np.float32))
grad = Tensor(np.random.randn(2, 2).astype(np.float32))
i = Tensor(1, mstype.int32)
grad_net = C.grad_all_with_sens(net)
grad_net(i, input1, grad)
def test_expand_dims_grad():
""" test_expand_dims_grad """
input_tensor = Tensor(np.array([[2, 2], [2, 2]]))
expand_dims = P.ExpandDims()
def fn(x):
output = expand_dims(x, 0)
return output
out = fn(input_tensor)
gfn = grad_all_with_sens(fn)
sens = Tensor(np.ones_like(out.asnumpy()))
args = [input_tensor, sens]
gout = gfn(*args)
expect = np.ones([2, 2])
assert np.all(gout[0].asnumpy() == expect)
def test_net_vargs_expand():
class AddNet(Cell):
def __init__(self):
super(AddNet, self).__init__()
self.w = Parameter(Tensor(np.ones((3, 4, 5), np.float32)),
"w2",
requires_grad=True)
def construct(self, x, y):
return x + y
x = Tensor(np.random.normal(0, 1, [3, 4, 5]).astype(np.float32))
y = Tensor(np.random.normal(0, 1, [3, 4, 5]).astype(np.float32))
sens = Tensor(np.random.normal(0, 1, [3, 4, 5]).astype(np.float32))
net = AddNet()
_ = C.grad_all_with_sens(net, net.trainable_params())(x, y, sens)
def test_squeeze_grad():
""" test_squeeze_grad """
input_tensor = Tensor(np.ones(shape=[3, 2, 1]))
squeeze = P.Squeeze(2)
def fn(x):
output = squeeze(x)
return output
out = fn(input_tensor)
gfn = grad_all_with_sens(fn)
sens = Tensor(np.ones_like(out.asnumpy()))
args = [input_tensor, sens]
gout = gfn(*args)
expect = np.ones([3, 2, 1])
assert np.all(gout[0].asnumpy() == expect)
def test_transpose_grad():
""" test_transpose_grad """
input_tensor = Tensor(np.array([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]))
perm = (0, 2, 1)
transpose = P.Transpose()
def fn(x):
output = transpose(x, perm)
return output
out = fn(input_tensor)
gfn = grad_all_with_sens(fn)
sens = Tensor(np.ones_like(out.asnumpy()))
args = [input_tensor, sens]
gout = gfn(*args)
expect = np.ones([2, 2, 3])
assert np.all(gout[0].asnumpy() == expect)
def test_reshape_grad():
""" test_reshape_grad """
input_tensor = Tensor(np.array([[-0.1, 0.3, 3.6], [0.4, 0.5, -3.2]]))
shp = (3, 2)
reshape = P.Reshape()
def fn(x):
output = reshape(x, shp)
return output
out = fn(input_tensor)
gfn = grad_all_with_sens(fn)
sens = Tensor(np.ones_like(out.asnumpy()))
args = [input_tensor, sens]
gout = gfn(*args)
expect = np.ones([2, 3])
assert np.all(gout[0].asnumpy() == expect)
def test_cast_grad():
""" test_cast_grad """
input_np = np.random.randn(2, 3).astype(np.float32)
input_x = Tensor(input_np)
td = ms.int32
cast = P.Cast()
def fn(x):
output = cast(x, td)
return output
out = fn(input_x)
gfn = grad_all_with_sens(fn)
sens = Tensor(np.ones_like(out.asnumpy()))
args = [input_x, sens]
gout = gfn(*args)
expect = np.ones((2, 3), dtype=np.float32)
assert np.all(gout[0].asnumpy() == expect)
def test_SubGrad():
""" test_SubGrad """
input_x = Tensor(np.array([[2, 2]]))
input_y = Tensor(np.array([[2, 2], [2, 2]]))
sub = P.Sub()
def fn(x, y):
output = sub(x, y)
return output
out = fn(input_x, input_y)
gfn = grad_all_with_sens(fn)
sens = Tensor(np.ones_like(out.asnumpy()))
args = [input_x, input_y, sens]
gout = gfn(*args)
expect_dx = np.ones([1, 2]).astype(np.int32) * 2 # reduce sum dout to the shape of x
expect_dy = np.ones([2, 2]).astype(np.int32) * (-1)
assert np.array_equal(gout[0].asnumpy(), expect_dx)
assert np.array_equal(gout[1].asnumpy(), expect_dy)
def test_parser_switch_layer_inputs_tuple():
class TwoInputTupleFinalNet(nn.Cell):
def __init__(self, funcs):
super().__init__()
self.funcs = funcs
def construct(self, i, inputa, inputb):
inputs = (inputa, inputb)
x = self.funcs[i](inputs)
return x
class Add(nn.Cell):
def __init__(self):
super().__init__()
self.op = P.TensorAdd()
def construct(self, x):
y = self.op(x[0], x[1])
return self.op(x[0], y)
class Mul(nn.Cell):
def __init__(self):
super().__init__()
self.op = P.Mul()
def construct(self, x):
y = self.op(x[0], x[1])
return self.op(x[0], y)
func1 = Add()
func2 = Mul()
funcs = (func1, func2)
net = TwoInputTupleFinalNet(funcs)
input1 = Tensor(np.random.randn(2, 3, 4, 5).astype(np.float32))
input2 = Tensor(np.random.randn(2, 3, 4, 5).astype(np.float32))
i = Tensor(1, mstype.int32)
grad = Tensor(np.random.randn(2, 3, 4, 5).astype(np.float32))
back_net = C.grad_all_with_sens(net)
back_out = back_net(i, input1, input2, grad)
def test_MulGrad():
""" test_MulGrad """
input_x = Tensor(np.array([[2, 2], [2, 2]], np.float32))
input_y = Tensor(np.array([[3, 3], [3, 3]], np.float32))
mul = P.Mul()
def fn(x, y):
output = mul(x, y)
return output
out = fn(input_x, input_y)
gfn = grad_all_with_sens(fn)
sens = Tensor(np.ones_like(out.asnumpy()) * 3)
args = [input_x, input_y, sens]
gout = gfn(*args)
expect_dx = np.ones([2, 2], np.float32) * 9
expect_dy = np.ones([2, 2], np.float32) * 6
assert np.all(gout[0].asnumpy().shape == expect_dx.shape)
assert np.all(gout[0].asnumpy() == expect_dx)
assert np.all(gout[1].asnumpy().shape == expect_dy.shape)
assert np.all(gout[1].asnumpy() == expect_dy)
def test_select_grad():
""" test_select_grad """
select = P.Select()
cond = Tensor(np.array([[True, False, False], [False, True, True]]))
x = Tensor(np.array([[1, 2, 3], [4, 5, 6]]).astype(np.float32))
y = Tensor(np.array([[7, 8, 9], [10, 11, 12]]).astype(np.float32))
def fn(cond, x, y):
output = select(cond, x, y)
return output
out = fn(cond, x, y)
gfn = grad_all_with_sens(fn)
sens = Tensor(np.ones_like(out.asnumpy()).astype(np.float32))
args = [cond, x, y, sens]
gout = gfn(*args)
expect_cond = np.zeros_like(cond)
expect_x = np.array([[1, 0, 0], [0, 1, 1]])
expect_y = np.array([[0, 1, 1], [1, 0, 0]])
assert np.all(gout[0].asnumpy() == expect_cond)
assert np.all(gout[1].asnumpy() == expect_x)
assert np.all(gout[2].asnumpy() == expect_y)
def construct(self, x):
return C.grad_all_with_sens(self.network)(x, self.sens)
def test_grad_multi_outputs():
assert C.grad_all_with_sens(multi_outputs)(2, 3, (1, 1)) == (4, 4)
def construct(self, *inputs):
return C.grad_all_with_sens(self.net)(*inputs)
def construct(self, x, y, b):
loss = self.network(x, y, b)
sens = P.Fill()(mstype.float32, P.Shape()(loss), 1.0)
return C.grad_all_with_sens(self.network)(x, y, b, sens)
def construct(self, x, y, z, sens):
return grad_all_with_sens(self.network)(x, y, z, sens)
def construct(self, inputa, inputb, inputz, output_grad):
gout = grad_all_with_sens(self.network)(inputa, inputb, inputz,
output_grad)
return gout
def first_derivative_dual(x, y):
""" first_derivative_dual """
return grad_all_with_sens(dual)(x, y, 1)
def construct(self, x, y, sens):
return C.grad_all_with_sens(self.net)(x, y, sens)
def df_multi_outputs(x, y):
return C.grad_all_with_sens(multi_outputs)(x, y, (1, 1))
def construct(self, input,z, w, output_grad):
return grad_all_with_sens(self.network)(input,z,w, output_grad)
def construct(self, x, y, output_grad):
return grad_all_with_sens(self.network)(x, y, output_grad)
def third_derivative_dual(x, y):
""" third_derivative_dual """
grad_fn = grad_all_with_sens(second_derivative_dual)
dfdx = grad_fn(x, y, (1, 0))[0]
dfdy = grad_fn(x, y, (0, 1))[1]
return dfdx, dfdy
