| |
| |
| |
| |
| |
| |
|
|
| import torch |
|
|
| def split_chessboard(x, num_split): |
| """ |
| x: b * c * h * w |
| Deividing x into num_split**2 sub-squares, and concatenate all the sub-squares on the batch dimension |
| """ |
| B, C, H, W = x.shape |
| assert H % num_split == 0 and W % num_split == 0 |
| h, w = H // num_split, W // num_split |
| x_split = torch.cat([x[:, :, i*h:(i+1)*h, j*w:(j+1)*w] for i in range(num_split) for j in range(num_split)], dim=0) |
| return x_split |
|
|
| def merge_chessboard(x, num_split): |
| """ |
| x: b * c * h * w |
| Assuming x contains num_split**2 sub-squares concatenated along batch dimension, merge the sub-squares back to the original whole square. |
| (inverse of split_chessboard) |
| """ |
| B, C, H, W = x.shape |
| assert B % (num_split**2) == 0 |
| b = B // (num_split**2) |
| x_merge = torch.cat([torch.cat([x[(i*num_split + j)*b:(i*num_split + j + 1)*b] for j in range(num_split)], dim=-1) |
| for i in range(num_split)], dim=-2) |
| return x_merge |
|
|
| def batched_forward(model, x, batch_size=-1): |
| if batch_size == -1: |
| return model(x) |
| else: |
| x_batched = x.split(batch_size) |
| outs = [model(x) for x in x_batched] |
| return torch.cat(outs, dim=0) |
|
|
|
|
