|
import pytest |
|
import torch |
|
from ding.torch_utils.reshape_helper import fold_batch, unfold_batch, unsqueeze_repeat |
|
|
|
|
|
@pytest.mark.unittest |
|
def test_fold_unfold_batch(): |
|
T, B, C, H, W = 10, 20, 3, 255, 255 |
|
data = torch.randn(T, B, C, H, W) |
|
data, batch_dim = fold_batch(data, nonbatch_ndims=3) |
|
assert data.shape == (T * B, C, H, W) and batch_dim == (T, B) |
|
data = unfold_batch(data, batch_dim) |
|
assert data.shape == (T, B, C, H, W) |
|
|
|
T, B, N = 10, 20, 100 |
|
data = torch.randn(T, B, N) |
|
data, batch_dim = fold_batch(data, nonbatch_ndims=1) |
|
assert data.shape == (T * B, N) and batch_dim == (T, B) |
|
data = unfold_batch(data, batch_dim) |
|
assert data.shape == (T, B, N) |
|
|
|
|
|
@pytest.mark.unittest |
|
def test_unsqueeze_repeat(): |
|
T, B, C, H, W = 10, 20, 3, 255, 255 |
|
repeat_times = 4 |
|
data = torch.randn(T, B, C, H, W) |
|
ensembled_data = unsqueeze_repeat(data, repeat_times) |
|
assert ensembled_data.shape == (repeat_times, T, B, C, H, W) |
|
|
|
ensembled_data = unsqueeze_repeat(data, repeat_times, -1) |
|
assert ensembled_data.shape == (T, B, C, H, W, repeat_times) |
|
|
|
ensembled_data = unsqueeze_repeat(data, repeat_times, 2) |
|
assert ensembled_data.shape == (T, B, repeat_times, C, H, W) |
|
|
