File size: 908 Bytes
36c95ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
from time import time
import pytest
import torch
import kornia
shapes = [(512, 3, 256, 256), (256, 1, 64, 64)]
PSs = [224, 32]
@pytest.mark.xfail(reason='May cause memory issues.')
def test_performance_speed(device, dtype):
if device.type != 'cuda' or not torch.cuda.is_available():
pytest.skip("Cuda not available in system,")
print("Benchmarking warp_affine")
for input_shape in shapes:
for PS in PSs:
BS = input_shape[0]
inpt = torch.rand(input_shape).to(device)
As = torch.eye(3).unsqueeze(0).repeat(BS, 1, 1)[:, :2, :].to(device)
As += 0.1 * torch.rand(As.size()).to(device)
torch.cuda.synchronize(device)
t = time()
_ = kornia.warp_affine(inpt, As, (PS, PS))
print(f"inp={input_shape}, PS={PS}, dev={device}, {time() - t}, sec")
torch.cuda.synchronize(device)
|