def test_slicing(test_case):
size = flow.Size([2, 3, 4, 5])
test_case.assertTrue(size[1:3] == flow.Size((3, 4)))
test_case.assertTrue(size[1:] == flow.Size((3, 4, 5)))
test_case.assertTrue(size[:2] == (2, 3))
test_case.assertTrue(size[-3:] == flow.Size((3, 4, 5)))
test_case.assertTrue(size[-3:-1] == flow.Size((3, 4)))
def _test_flatten(test_case, device):
m = flow.nn.Flatten()
x = flow.Tensor(32, 2, 5, 5, device=flow.device(device))
flow.nn.init.uniform_(x)
y = m(x)
test_case.assertTrue(y.shape == flow.Size((32, 50)))
test_case.assertTrue(
np.array_equal(y.numpy().flatten(),
x.numpy().flatten()))
y2 = flow.flatten(x, start_dim=2)
test_case.assertTrue(y2.shape == flow.Size((32, 2, 25)))
test_case.assertTrue(
np.array_equal(y2.numpy().flatten(),
x.numpy().flatten()))
y3 = x.flatten(start_dim=1)
test_case.assertTrue(y3.shape == flow.Size((32, 50)))
test_case.assertTrue(
np.array_equal(y3.numpy().flatten(),
x.numpy().flatten()))
y4 = x.flatten(start_dim=1, end_dim=2)
test_case.assertTrue(y4.shape == flow.Size((32, 10, 5)))
test_case.assertTrue(
np.array_equal(y4.numpy().flatten(),
x.numpy().flatten()))
y5 = flow.flatten(x)
test_case.assertTrue(y5.shape == flow.Size((1600, )))
test_case.assertTrue(
np.array_equal(y5.numpy().flatten(),
x.numpy().flatten()))
def test_affine_grid_3d(test_case):
input = flow.tensor(np.arange(1.0, 13).reshape((1, 3, 4)),
dtype=flow.float32)
output = flow.nn.functional.affine_grid(input,
flow.Size([1, 1, 2, 2, 2]),
align_corners=True)
groundtruth = np.array([[
[
[[-2.0, -10.0, -18.0], [0.0, 0.0, 0.0]],
[[2.0, 2.0, 2.0], [4.0, 12.0, 20.0]],
],
[
[[4.0, 4.0, 4.0], [6.0, 14.0, 22.0]],
[[8.0, 16.0, 24.0], [10.0, 26.0, 42.0]],
],
]])
test_case.assertTrue(
np.allclose(output.numpy(), groundtruth, rtol=1e-3, atol=1e-4))
output = flow.nn.functional.affine_grid(input,
flow.Size([1, 1, 2, 2, 2]),
align_corners=False)
groundtruth = np.array([[
[
[[1.0, -1.0, -3.0], [2.0, 4.0, 6.0]],
[[3.0, 5.0, 7.0], [4.0, 10.0, 16.0]],
],
[
[[4.0, 6.0, 8.0], [5.0, 11.0, 17.0]],
[[6.0, 12.0, 18.0], [7.0, 17.0, 27.0]],
],
]])
test_case.assertTrue(
np.allclose(output.numpy(), groundtruth, rtol=1e-3, atol=1e-4))
def test_equal(test_case):
size = flow.Size((2, 3))
test_case.assertEqual(size == (2, 3), True)
test_case.assertEqual(size == (3, 2), False)
test_case.assertEqual(size == flow.Size((2, 3)), True)
test_case.assertEqual(size == flow.Size((3, 2)), False)
test_case.assertEqual(size == [2, 3], False)
test_case.assertEqual(size == dict(), False)
def test_dynamic_attrs(test_case):
x = (flow.builtin_op("constant").Output("out").Attr(
"is_floating_value",
True).Attr("floating_value",
3.0).Attr("dtype",
flow.float32).Attr("shape",
[2, 3]).Build())()[0]
op = flow.builtin_op("expand_dims").Input("in").Output("out").Build()
y = op(x, axis=1)[0]
test_case.assertEqual(y.shape, flow.Size((2, 1, 3)))
y = op(x, axis=2)[0]
test_case.assertEqual(y.shape, flow.Size((2, 3, 1)))
def test_parital_fc(test_case):
p = flow.env.all_device_placement("cuda")
w = flow.randn(50000, 128, placement=p, sbp=flow.sbp.broadcast)
label = flow.randint(0,
50000, (512, ),
placement=p,
sbp=flow.sbp.broadcast)
num_sample = 5000
out = flow.distributed_partial_fc_sample(w, label, num_sample)
test_case.assertTrue(out[0].shape == flow.Size([512]))
test_case.assertTrue(out[1].shape == flow.Size([5000]))
test_case.assertTrue(out[2].shape == flow.Size([5000, 128]))
def test_parital_fc(test_case):
p = flow.env.all_device_placement("cuda")
w = flow.randn(
50000, 128, placement=p, sbp=flow.sbp.broadcast, requires_grad=True
)
label = flow.randint(0, 50000, (512,), placement=p, sbp=flow.sbp.broadcast)
num_sample = 5000
out = flow.distributed_partial_fc_sample(w, label, num_sample)
test_case.assertTrue(out[0].shape == flow.Size([512]))
test_case.assertTrue(out[1].shape == flow.Size([5000]))
test_case.assertTrue(out[2].shape == flow.Size([5000, 128]))
# test gradient function
sample_weight = out[2]
sample_weight.sum().backward()
def test_size(test_case):
size = flow.Size((4, 3, 10, 5))
test_case.assertTrue(size[0] == 4)
test_case.assertTrue(size[2] == 10)
test_case.assertTrue(len(size) == 4)
size = flow.Size([4, 3, 10, 5])
test_case.assertTrue(size[0] == 4)
test_case.assertTrue(size[2] == 10)
test_case.assertTrue(len(size) == 4)
size = flow.Size(size)
test_case.assertTrue(size[0] == 4)
test_case.assertTrue(size[2] == 10)
test_case.assertTrue(len(size) == 4)
def _test_consistent_randint(test_case, shape, placement, sbp, dtype):
x = flow.randint(1, 10, shape, placement=placement, sbp=sbp, dtype=dtype)
test_case.assertEqual(x.shape, flow.Size(shape))
test_case.assertEqual(x.sbp, sbp)
test_case.assertEqual(x.placement, placement)
test_case.assertEqual(x.dtype, dtype)
def _test_graph_empty(test_case, func, shape, placement, sbp):
func2 = None
if func == "empty":
func = flow.empty
elif func == "new_empty":
func = flow.empty
func2 = flow.new_empty
else:
raise NotImplementedError
class ConsistentEmptyGraph(flow.nn.Graph):
def __init__(self, ):
super().__init__()
def build(self):
x = func(*shape, placement=placement, sbp=sbp)
if func2:
x = func2(x, size=shape)
return x
model = ConsistentEmptyGraph()
x = model()
test_case.assertEqual(x.shape, flow.Size(shape))
test_case.assertEqual(x.sbp, sbp)
test_case.assertEqual(x.placement, placement)
def _compare_with_np(
test_case, x_shape, dtype,
):
x = np.random.randn(*x_shape).astype(type_name_to_np_type[dtype])
ret = flow.Size(x_shape)
for idx in range(0, len(ret)):
test_case.assertEqual(ret[idx], x.shape[idx])
def _test_slice_empty(test_case, device):
np_arr = np.random.randn(10).astype(np.float32)
x = flow.tensor(np_arr, device=flow.device(device))
y = x[3:3]
test_case.assertTrue(y.shape, flow.Size((0, )))
np_out = np_arr[3:3]
test_case.assertTrue(np.array_equal(y.numpy(), np_out))
def _test_consistent_constant(test_case, func, shape, placement, sbp):
func2 = None
if func == "ones":
func = flow.ones
np_res = np.ones(shape)
elif func == "zeros":
func = flow.zeros
np_res = np.zeros(shape)
elif func == "new_zeros":
func = flow.zeros
np_res = np.zeros(shape)
func2 = flow.new_zeros
else:
raise NotImplementedError
x = func(*shape, placement=placement, sbp=sbp)
if func2:
x = func2(x)
test_case.assertEqual(x.shape, flow.Size(shape))
test_case.assertEqual(x.sbp, sbp)
test_case.assertEqual(x.placement, placement)
x = x.to_global(
placement=placement,
sbp=[flow.sbp.broadcast for _ in range(len(placement.ranks.shape))],
).to_local()
test_case.assertTrue(np.array_equal(x.numpy(), np_res))
def test_affine_grid_2d(test_case):
input = flow.tensor(np.arange(1.0, 7).reshape((1, 2, 3)), dtype=flow.float32)
output = flow.nn.functional.affine_grid(
input, flow.Size([1, 1, 2, 2]), align_corners=True
)
groundtruth = np.array([[[[0.0, -3.0], [2.0, 5.0]], [[4.0, 7.0], [6.0, 15.0]]]])
test_case.assertTrue(
np.allclose(output.numpy(), groundtruth, rtol=1e-3, atol=1e-4)
)
output = flow.nn.functional.affine_grid(
input, flow.Size([1, 1, 2, 2]), align_corners=False
)
groundtruth = np.array([[[[1.5, 1.5], [2.5, 5.5]], [[3.5, 6.5], [4.5, 10.5]]]])
test_case.assertTrue(
np.allclose(output.numpy(), groundtruth, rtol=1e-3, atol=1e-4)
)
def _test_consistent_randn(test_case, shape, placement, sbp):
x1 = flow.randn(*shape, placement=placement, sbp=sbp)
x2 = flow.randn(*shape, placement=placement, sbp=sbp)
test_case.assertTrue(
not np.allclose(x1.numpy(), x2.numpy(), atol=1e-4, rtol=1e-4))
test_case.assertEqual(x1.shape, flow.Size(shape))
test_case.assertEqual(x1.sbp, sbp)
test_case.assertEqual(x1.placement, placement)
def test_mirrored_tensor_and_op(test_case):
x1 = flow.Tensor([[1.0, 2.0]])
test_case.assertEqual(x1.dtype, flow.float32)
test_case.assertEqual(x1.shape, flow.Size((1, 2)))
x2 = flow.Tensor([[1.0], [2.0]])
y = flow.matmul(x1, x2)
test_case.assertTrue(
np.allclose(y.numpy(), np.array([[5.0]], dtype=np.float32)))
def _test_expand_flow_size(test_case, device):
input_shape = (2, 4, 1, 32)
expand_dim = flow.Size([2, 4, 2, 32])
input, gout, out_np, gin_np = _np_get_expand(input_shape, expand_dim)
of_input = flow.tensor(input, dtype=flow.int, device=flow.device(device))
of_out = of_input.expand(expand_dim)
test_case.assertTrue(
np.array_equal(of_out.numpy(), out_np.astype(np.int32)))
def _test_new_empty(test_case, shape, dtype, device, requires_grad):
x = flow.empty(shape, dtype=dtype, device=flow.device(device))
y = x.new_empty(
shape,
dtype=dtype,
device=flow.device(device),
requires_grad=requires_grad if dtype == flow.float32 else False,
)
test_case.assertFalse(y.is_global)
test_case.assertEqual(y.shape, flow.Size(shape))
test_case.assertEqual(y.dtype, dtype)
test_case.assertEqual(y.device, flow.device(device))
if dtype == flow.float32:
test_case.assertEqual(y.requires_grad, requires_grad)
y = x.new_empty(*shape)
test_case.assertFalse(y.is_global)
test_case.assertEqual(y.shape, flow.Size(shape))
test_case.assertEqual(y.dtype, x.dtype)
test_case.assertEqual(y.device, x.device)
test_case.assertFalse(y.requires_grad)
def _test_local_empty(test_case, shape, dtype, device, requires_grad):
x = flow.empty(
shape,
dtype=dtype,
device=flow.device(device),
requires_grad=requires_grad if dtype == flow.float32 else False,
)
test_case.assertFalse(x.is_consistent)
test_case.assertEqual(x.shape, flow.Size(shape))
test_case.assertEqual(x.dtype, dtype)
test_case.assertEqual(x.device, flow.device(device))
if dtype == flow.float32:
test_case.assertEqual(x.requires_grad, requires_grad)
def test_index(test_case):
size = flow.Size((2, 3, 2, 4, 4))
test_case.assertEqual(size.index(2), 0)
test_case.assertEqual(size.index(2, start=0), 0)
test_case.assertEqual(size.index(2, start=0, end=20), 0)
test_case.assertEqual(size.index(2, start=1, end=20), 2)
test_case.assertEqual(size.index(4), 3)
test_case.assertEqual(size.index(4, start=4), 4)
with test_case.assertRaises(ValueError):
size.index(4, start=0, end=3)
with test_case.assertRaises(ValueError):
size.index(5)
with test_case.assertRaises(ValueError):
size.index(2, start=3)
def _test_graph_rand(test_case, shape, placement, sbp):
class ConsistentRandGraph(flow.nn.Graph):
def __init__(self, ):
super().__init__()
def build(self):
x = flow.rand(*shape, placement=placement, sbp=sbp)
return x
model = ConsistentRandGraph()
x = model()
test_case.assertEqual(x.shape, flow.Size(shape))
test_case.assertEqual(x.sbp, sbp)
test_case.assertEqual(x.placement, placement)
def _test_consistent_empty(test_case, shape, dtype, placement, sbp, requires_grad):
placement = flow.placement(placement, {0: [0]})
x = flow.empty(
shape,
dtype=dtype,
placement=placement,
sbp=sbp,
requires_grad=requires_grad if dtype == flow.float32 else False,
)
test_case.assertTrue(x.is_consistent)
test_case.assertEqual(x.shape, flow.Size(shape))
test_case.assertEqual(x.dtype, dtype)
test_case.assertEqual(x.placement, placement)
test_case.assertEqual(x.sbp[0], sbp)
if dtype == flow.float32:
test_case.assertEqual(x.requires_grad, requires_grad)
def _test_consistent_empty(test_case, func, shape, placement, sbp):
func2 = None
if func == "empty":
func = flow.empty
elif func == "new_empty":
func = flow.empty
func2 = flow.new_empty
else:
raise NotImplementedError
x = func(*shape, placement=placement, sbp=sbp)
if func2:
x = func2(x, size=shape)
test_case.assertEqual(x.shape, flow.Size(shape))
test_case.assertEqual(x.sbp, sbp)
test_case.assertEqual(x.placement, placement)
def test_moduledict(test_case):
class ModuleDict(nn.Module):
def __init__(self):
super(ModuleDict, self).__init__()
self.choices = nn.ModuleDict(
{"conv": nn.Conv2d(10, 10, 3), "pool": nn.MaxPool2d(3)}
)
self.activations = nn.ModuleDict(
{"relu": nn.ReLU(), "prelu": nn.PReLU()}
)
def forward(self, x, choice, act):
x = self.choices[choice](x)
x = self.activations[act](x)
return x
model = ModuleDict()
input = flow.tensor(np.random.randn(4, 10, 32, 32), dtype=flow.float32)
output = model(input, "conv", "relu")
test_case.assertEqual(output.shape, flow.Size([4, 10, 30, 30]))
def _test_view_flow_size(test_case, device):
x = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12],
[13, 14, 15, 16]]).astype(np.float32)
input = flow.tensor(x,
dtype=flow.float32,
device=flow.device(device),
requires_grad=True)
shape = flow.Size([2, 2, 2, -1])
of_out = input.view(shape)
np_shape = (2, 2, 2, 2)
test_case.assertTrue(np.array_equal(of_out.numpy().shape, np_shape))
of_out = of_out.sum()
of_out.backward()
np_grad = np.array([
[1.0, 1.0, 1.0, 1.0],
[1.0, 1.0, 1.0, 1.0],
[1.0, 1.0, 1.0, 1.0],
[1.0, 1.0, 1.0, 1.0],
])
test_case.assertTrue(
np.allclose(np_grad, input.grad.numpy(), 0.0001, 0.0001))
def _test_graph_constant(test_case, func, shape, placement, sbp):
func2 = None
if func == "ones":
func = flow.ones
np_res = np.ones(shape)
elif func == "zeros":
func = flow.zeros
np_res = np.zeros(shape)
elif func == "new_zeros":
func = flow.zeros
np_res = np.zeros(shape)
func2 = flow.new_zeros
else:
raise NotImplementedError
class ConsistentConstantGraph(flow.nn.Graph):
def __init__(self, ):
super().__init__()
def build(self):
x = func(*shape, placement=placement, sbp=sbp)
if func2:
x = func2(x)
return x
model = ConsistentConstantGraph()
x = model()
test_case.assertEqual(x.shape, flow.Size(shape))
test_case.assertEqual(x.sbp, sbp)
test_case.assertEqual(x.placement, placement)
x = x.to_global(
placement=placement,
sbp=[flow.sbp.broadcast for _ in range(len(placement.ranks.shape))],
).to_local()
test_case.assertTrue(np.array_equal(x.numpy(), np_res))
def _test_consistent_rand(test_case, shape, placement, sbp):
x = flow.rand(*shape, placement=placement, sbp=sbp)
test_case.assertEqual(x.shape, flow.Size(shape))
test_case.assertEqual(x.sbp, sbp)
test_case.assertEqual(x.placement, placement)
def _test_randn_with_flow_size(test_case, device, shape):
y1 = flow.randn(flow.Size(shape), device=flow.device(device))
y2 = flow.randn(flow.Size(shape), device=flow.device(device))
test_case.assertTrue(not np.array_equal(y1.numpy(), y2.numpy()))
test_case.assertTrue(shape == y1.shape)
def test_unpack(test_case):
one, two, three, four = flow.Size((1, 2, 3, 4))
test_case.assertEqual(one, 1)
test_case.assertEqual(two, 2)
test_case.assertEqual(three, 3)
test_case.assertEqual(four, 4)
def test_count(test_case):
size = flow.Size((2, 2, 3, 4))
test_case.assertEqual(size.count(1), 0)
test_case.assertEqual(size.count(2), 2)
test_case.assertEqual(size.count(3), 1)
test_case.assertEqual(size.count(4), 1)