diff --git a/test/variable_size/test_sequence.py b/test/variable_size/test_sequence.py new file mode 100644 index 0000000000000000000000000000000000000000..3fea0ea4000ed1b698292b9b1260a0a51f5da376 --- /dev/null +++ b/test/variable_size/test_sequence.py @@ -0,0 +1,91 @@ +import time +import numpy as np +import torch +import nestedtensor +from ding.torch_utils import Transformer + + +def generate_variable_sequence(N, M, sample_range): + return [torch.zeros(1, np.random.randint(*sample_range), M) for _ in range(N)] + + +def naive_method(model, data, cuda, test_loop=3): + result = [] + for _ in range(test_loop): + t_start = time.time() + with torch.no_grad(): + output = [model(d) for d in data] + if cuda: + torch.cuda.synchronize() + t_end = time.time() + result.append(t_end - t_start) + print('naive_method test time avg: {}, max: {}'.format(np.mean(result), np.max(result))) + return output, result + + +def padding_method(model, data, cuda, test_loop=3): + result = [] + for _ in range(test_loop): + t_start = time.time() + with torch.no_grad(): + # padding + max_n = max([d.shape[1] for d in data]) + new_data = torch.zeros(len(data), max_n, data[0].shape[-1]).to(data[0].device) + mask = torch.zeros(len(data), max_n) + for i in range(len(data)): + mask[i, :data[i].shape[1]].add_(1) + mask = mask.bool().to(data[0].device) + + padding_output = model(new_data, mask=mask) + + output = [] + for i in range(len(data)): + output.append(padding_output[i, :data[i].shape[1]].unsqueeze(0)) + if cuda: + torch.cuda.synchronize() + t_end = time.time() + result.append(t_end - t_start) + print('padding_method test time avg: {}, max: {}'.format(np.mean(result), np.max(result))) + return output, result + + +def nestedtensor_method(model, data, cuda, test_loop=3): + raise NotImplementedError("nestedtensor doesn't support chunk op now") + result = [] + data = nestedtensor.nested_tensor([d.squeeze(0) for d in data]) + for _ in range(test_loop): + t_start = time.time() + with torch.no_grad(): + output = model(data) + if cuda: + torch.cuda.synchronize() + t_end = time.time() + result.append(t_end - t_start) + print('nestedtensor_method test time avg: {}, max: {}'.format(np.mean(result), np.max(result))) + output = [o.unsqueeze(0) for o in output] + return output, result + + +def main(cuda): + N, M = 64, 128 + sample_range = [32, 64] + np.random.seed(0) + + data = generate_variable_sequence(N, M, sample_range) + model = Transformer(input_dim=M) + print(model) + if cuda: + model.cuda() + data = [d.cuda() for d in data] + # warm up + for _ in range(10): + with torch.no_grad(): + model(data[0]) + + naive_output, naive_result = naive_method(model, data, cuda) + padding_output, padding_result = padding_method(model, data, cuda) + # nest_output, nest_result = nestedtensor_method(model, data, cuda) + + +if __name__ == "__main__": + main(cuda=False) diff --git a/test/variable_size/test_unet.py b/test/variable_size/test_unet.py index 5bffd2cf0f187133d913cdfcdc045692274c749b..56dbfdb525b0d983f7fa414140733a6b222066f8 100644 --- a/test/variable_size/test_unet.py +++ b/test/variable_size/test_unet.py @@ -34,6 +34,36 @@ def naive_method(model, data, cuda, test_loop=3): return output, result +def padding_method(model, data, cuda, test_loop=3): + result = [] + for _ in range(test_loop): + t_start = time.time() + with torch.no_grad(): + # padding + max_h = max([d.shape[-2] for d in data]) + max_w = max([d.shape[-1] for d in data]) + new_data = torch.zeros(len(data), 3, max_h, max_w).to(data[0].device) + start_h = [max_h - d.shape[-2] for d in data] + start_w = [max_w - d.shape[-1] for d in data] + for i in range(len(data)): + new_data[i, :, start_h[i]:, start_w[i]:] = data[i] + + padding_output = model(new_data) + + output = [] + for i in range(len(data)): + output.append( + padding_output[i, :, start_h[i]:start_h[i] + data[i].shape[-2], + start_w[i]:start_w[i] + data[i].shape[-1]] + ) + if cuda: + torch.cuda.synchronize() + t_end = time.time() + result.append(t_end - t_start) + print('padding_method test time avg: {}, max: {}'.format(np.mean(result), np.max(result))) + return output, result + + def nestedtensor_method(model, data, cuda, test_loop=3): result = [] data = nestedtensor.nested_tensor([d.squeeze(0) for d in data]) @@ -74,7 +104,8 @@ def main(cuda): same_output, same_result = same_test(model, data, cuda) naive_output, naive_result = naive_method(model, data, cuda) - assert len(naive_output) == B + assert len(naive_output) == B, len(naive_output) + padding_output, padding_result = padding_method(model, data, cuda) nest_output, nest_result = nestedtensor_method(model, data, cuda) print(naive_output[0][0, 0, 0, :10]) print(nest_output[0][0, 0, 0, :10])