| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class Correlation(nn.Module): |
| def __init__(self, max_displacement=4, *args, **kwargs): |
| super(Correlation, self).__init__() |
| self.max_displacement = max_displacement |
| self.output_dim = 2 * self.max_displacement + 1 |
| self.pad_size = self.max_displacement |
|
|
| def forward(self, x1, x2): |
| B, C, H, W = x1.size() |
|
|
| x2 = F.pad(x2, [self.pad_size] * 4) |
| cv = [] |
| for i in range(self.output_dim): |
| for j in range(self.output_dim): |
| cost = x1 * x2[:, :, i : (i + H), j : (j + W)] |
| cost = torch.mean(cost, 1, keepdim=True) |
| cv.append(cost) |
| return torch.cat(cv, 1) |
|
|
|
|
| if __name__ == "__main__": |
| import random |
| import time |
|
|
| from correlation_package.correlation import Correlation as Correlation_cuda |
|
|
| device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") |
| corr1 = Correlation( |
| max_displacement=4, kernel_size=1, stride1=1, stride2=1, corr_multiply=1 |
| ).to(device) |
|
|
| corr2 = Correlation_cuda( |
| pad_size=4, |
| kernel_size=1, |
| max_displacement=4, |
| stride1=1, |
| stride2=1, |
| corr_multiply=1, |
| ) |
|
|
| t1_sum = 0 |
| t2_sum = 0 |
|
|
| for i in range(50): |
| C = random.choice([128, 256]) |
| H = random.choice([128, 256]) |
| W = random.choice([64, 128]) |
| x1 = torch.randn(4, C, H, W, requires_grad=True).to(device) |
| x2 = torch.randn(4, C, H, W).to(device) |
|
|
| end = time.time() |
| y2 = corr2(x1, x2) |
| t2_f = time.time() - end |
|
|
| end = time.time() |
| y2.sum().backward() |
| t2_b = time.time() - end |
|
|
| end = time.time() |
| y1 = corr1(x1, x2) |
| t1_f = time.time() - end |
|
|
| end = time.time() |
| y1.sum().backward() |
| t1_b = time.time() - end |
|
|
| assert torch.allclose(y1, y2, atol=1e-7) |
|
|
| print( |
| "Forward: cuda: {:.3f}ms, pytorch: {:.3f}ms".format(t1_f * 100, t2_f * 100) |
| ) |
| print( |
| "Backward: cuda: {:.3f}ms, pytorch: {:.3f}ms".format(t1_b * 100, t2_b * 100) |
| ) |
|
|
| if i < 3: |
| continue |
| t1_sum += t1_b + t1_f |
| t2_sum += t2_b + t2_f |
|
|
| print("cuda: {:.3f}s, pytorch: {:.3f}s".format(t1_sum, t2_sum)) |
| ... |
|
|
