| from sudoku.loader import data_loader, DataIterBuffer, get_datasets, train_dataset | |
| from torch.utils.data import DataLoader | |
| def test_data_loader(): | |
| train_loader, test_loader = data_loader() | |
| X, Y = next(iter(train_loader)) | |
| assert X.shape == Y.shape == (32, 2, 9 * 9 * 9) | |
| assert (Y.sum(-1)[:, 0] == (9 * 9 * 8)).all() | |
| assert (Y.sum(-1)[:, 1] == (9 * 9)).all() | |
| assert (X.sum(-1)[:, 0] < (9 * 9 * 8)).all() | |
| assert (X.sum(-1)[:, 1] < (9 * 9)).all() | |
| def test_data_iter_buffer(): | |
| i = 0 | |
| data_iter = DataIterBuffer(train_dataset) | |
| data_loader_buffer = DataLoader(data_iter, 32) | |
| for X, Y in data_loader_buffer: | |
| i += 1 | |
| if i % 3 == 0: | |
| data_iter.append(X, Y) | |
| assert i == 29 | |
| def test_max_holes(): | |
| train_dataset, test_dataset = get_datasets(train_size=2, test_size=2, max_holes=2) | |
| X, Y = next(iter(train_dataset)) | |
| x_holes = X[1].reshape(9, 9, 9).sum(-1) == 0 | |
| print(X.shape, x_holes.shape) | |
| assert x_holes.sum() == 2 | |