def test_iloc(self):
torch.manual_seed(123)
a = [torch.randn((i, j)) for i, j in [(5, 3), (5, 2)]]
a = tuplefy(a, torch.randint(10, (5, )))
b = tuplefy((a[0][0][:2], a[0][1][:2]), a[1][:2])
a = a.iloc[:2]
assert_tupletree_equal(a, b)
def test_stack_shapes(self):
torch.manual_seed(123)
a = [torch.randn((i, j)) for i, j in [(5, 3), (5, 2)]]
a = tuplefy(a, torch.randint(10, (5, )))
b = tuplefy((torch.Size([3, 5, 3]), torch.Size([3, 5, 2])),
torch.Size([3, 5]))
assert_tupletree_equal(a.repeat(3).stack().shapes(), b)
def test_dataloader_input_only():
n = 10
input = tuplefy(torch.randn(n, 3), torch.randn(n, 2))
target = torch.randn(n)
dl = tuplefy(input, target).make_dataloader(4, False)
dl_input = dataloader_input_only(dl)
for (inp, tar), inp2 in zip(dl, dl_input):
assert_tupletree_equal(inp, inp2)
def test_make_data_loader_sorted(batch_size, num_workers):
torch.manual_seed(123)
a = tuplefy(torch.randn(10, 3), torch.randn(10))
dl = make_dataloader(a, batch_size, False, num_workers)
b = tuplefy()
for inp, tar in dl:
b = b + tuplefy(inp, tar).add_root()
assert_tupletree_equal(b.cat(), a)
def test_np2torch2np(self):
np.random.seed(123)
a = [
np.random.normal(size=(i, j)).astype('float32')
for i, j in [(2, 3), (1, 2)]
]
a = tuplefy(a, np.random.choice(10, size=(4, )))
b = a.to_tensor().to_numpy()
assert_tupletree_equal(a, b)
def test_next_iter(self, batch_size, num_workers):
torch.manual_seed(123)
n = 20
a = ((torch.randn(n, 3), torch.randint(200, (n, 2))), torch.randn(n))
a = tuplefy(a)
ds = DatasetTuple(*a)
dl = DataLoaderBatch(ds, batch_size, False, num_workers=num_workers)
a = a.iloc[:batch_size]
b = next(iter(dl))
assert_tupletree_equal(a, b)
def test_make_data_loader_unsorted(batch_size, num_workers):
torch.manual_seed(123)
a = tuplefy(torch.randn(10, 3), torch.randn(10))
dl = make_dataloader(a, batch_size, True, num_workers)
b = tuplefy()
for inp, tar in dl:
b = b + tuplefy(inp, tar).add_root()
b = b.cat()
assert_tupletree_equal(b.shapes(), a.shapes())
assert (tuplefy(a, b).apply(sum).zip_leaf().apply(
lambda x: x[0] - x[1]).apply(lambda x: x.max()).pipe(max) < 1e-6)
def test_step_with_model(self, optim_class):
torch.manual_seed(1234)
inp, tar = torch.randn(10, 3), torch.randn(10, 1)
net = torch.nn.Linear(3, 1)
model = Model(net, torch.nn.MSELoss(), optim_class(lr=0.1))
params = get_params(net)
model.fit(inp, tar, verbose=False)
new_params = get_params(net)
with pytest.raises(AssertionError) as ex:
assert_tupletree_equal(params, new_params)
assert str(ex.value) == "Not equal values"
def test_step(self, optimizer_class):
torch.manual_seed(1234)
net = torch.nn.Linear(5, 1)
optimizer = optimizer_class(net.parameters(), lr=0.1)
op = optim.OptimWrap(optimizer)
op.zero_grad()
params = get_params(net)
inp = torch.randn(10, 5)
net(inp).mean().backward()
op.step()
new_params = get_params(net)
with pytest.raises(AssertionError) as ex:
assert_tupletree_equal(params, new_params)
assert str(ex.value) == "Not equal values"
def test_getitem(self):
torch.manual_seed(123)
n = 10
a = ((torch.randn(n, 3), torch.randint(200, (n, 2))), torch.randn(n))
ds = DatasetTuple(*a)
assert_tupletree_equal(ds[0], ds[[0]])
assert_tupletree_equal(ds[0], ds[:1])
assert_tupletree_equal(ds[2:5], ds[[2, 3, 4]])
def test_assert_tupletree_equal_pass(a, b, check_dtypes):
assert_tupletree_equal(a, b, check_dtypes)
def test_assert_tupletree_equal_fails(a, b, ex_string):
with pytest.raises(AssertionError) as ex:
assert_tupletree_equal(a, b)
assert str(ex.value) == ex_string
def test_cat_split(self):
torch.manual_seed(123)
a = [torch.randn((i, j)) for i, j in [(5, 3), (5, 2)]]
a = tuplefy(a, torch.randint(10, (5, )))
for sub in a.repeat(3).cat().split(5):
assert_tupletree_equal(a, sub)
def test_torch2np2torch(self):
torch.manual_seed(123)
a = [torch.randn((i, j)) for i, j in [(2, 3), (1, 2)]]
a = tuplefy(a, torch.randint(10, (4, )))
b = a.to_numpy().to_tensor()
assert_tupletree_equal(a, b)
