MPM-PyTorch/benchmark_svd.py at master · wgsxm/MPM-PyTorch · GitHub

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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
import time
import torch
from tqdm import tqdm

from mpm_pytorch.constitutive_models.warp_svd import SVD

wp_svd = SVD()
wp_svd(torch.randn(1, 3, 3))
def warp_svd(x):
    return wp_svd(x)

def torch_svd(x):
    U, s, Vh = torch.svd(x)
    return U, s, Vh.transpose(-2, -1)

if __name__ == '__main__':
    test_time = 10
    n = 100000
    atol = 1e-5

    print("\nTest correctness")
    Is = torch.eye(3).unsqueeze(0).repeat(n, 1, 1)
    for _ in tqdm(range(test_time), desc="Test orthogonality of U and V"):
        x = torch.randn(n, 3, 3)
        U, s, V = warp_svd(x)
        UU = U @ U.transpose(-2, -1)
        VV = V @ V.transpose(-2, -1)
        assert torch.allclose(UU, Is, atol=1e-5)
        assert torch.allclose(VV, Is, atol=1e-5)
    bar = tqdm(range(test_time), desc="Test correctness of decomposition")
    for _ in bar:
        x = torch.randn(n, 3, 3)
        U, s, V = warp_svd(x)
        warp_pred = U @ torch.diag_embed(s) @ V
        warp_mae = torch.abs(warp_pred - x).mean()
        U, s, V = torch_svd(x)
        torch_pred = U @ torch.diag_embed(s) @ V
        torch_mae = torch.abs(torch_pred - x).mean()
        bar.set_postfix(warp_mae=warp_mae.item(), torch_mae=torch_mae.item())
        assert warp_mae < atol

    print("\nTest differentiability")
    bar = tqdm(range(test_time), desc="Test differentiability")
    for _ in bar:
        x = torch.randn(n, 3, 3, requires_grad=True)
        U, s, V = warp_svd(x)
        warp_pred = U @ torch.diag_embed(s) @ V
        warp_pred.sum().backward()
        assert x.grad is not None
        warp_mae = torch.abs(x.grad - torch.ones_like(x)).mean()
        x.grad = None
        U, s, V = torch_svd(x)
        torch_pred = U @ torch.diag_embed(s) @ V
        torch_pred.sum().backward()
        torch_mae = torch.abs(x.grad - torch.ones_like(x)).mean()
        bar.set_postfix(warp_mae=warp_mae.item(), torch_mae=torch_mae.item())
        assert warp_mae < atol

    print("\nTest speed")
    x = torch.randn(n, 3, 3)
    start = time.time()
    for _ in tqdm(range(test_time), desc="Test warp_svd forward (CPU)"):
        U, s, V = warp_svd(x)
    end = time.time()
    print("warp_svd: ", (end - start) / test_time, "seconds")
    start = time.time()
    for _ in tqdm(range(test_time), desc="Test torch_svd forward (CPU)"):
        U, s, V = torch_svd(x)
    end = time.time()
    print("torch_svd: ", (end - start) / test_time, "seconds")

    print()
    x = torch.randn(n, 3, 3, requires_grad=True)
    start = time.time()
    for _ in tqdm(range(test_time), desc="Test warp_svd forward + backward (CPU)"):
        U, s, V = warp_svd(x)
        warp_pred = U @ torch.diag_embed(s) @ V
        warp_pred.sum().backward()
    end = time.time()
    print("warp_svd: ", (end - start) / test_time, "seconds")
    x.grad = None
    start = time.time()
    for _ in tqdm(range(test_time), desc="Test torch_svd forward + backward (CPU)"):
        U, s, V = torch_svd(x)
        torch_pred = U @ torch.diag_embed(s) @ V
        torch_pred.sum().backward()
    end = time.time()
    print("torch_svd: ", (end - start) / test_time, "seconds")

    if torch.cuda.is_available():
        print()
        x = torch.randn(n, 3, 3).cuda().requires_grad_()
        start = time.time()
        for _ in tqdm(range(test_time), desc="Test warp_svd forward + backward (GPU)"):
            U, s, V = warp_svd(x)
            warp_pred = U @ torch.diag_embed(s) @ V
            warp_pred.sum().backward()
        end = time.time()
        print("warp_svd: ", (end - start) / test_time, "seconds")
        x.grad = None
        start = time.time()
        for _ in tqdm(range(test_time), desc="Test torch_svd forward + backward (GPU)"):
            U, s, V = torch_svd(x)
            torch_pred = U @ torch.diag_embed(s) @ V
            torch_pred.sum().backward()
        end = time.time()
        print("torch_svd: ", (end - start) / test_time, "seconds")