-
Notifications
You must be signed in to change notification settings - Fork 0
/
ChebConv.py
169 lines (139 loc) · 6.14 KB
/
ChebConv.py
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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
import torch
import torch.nn as nn
from torch.nn import init
import numpy as np
import scipy.sparse as sp
"""
@inproceedings{wu2021graph,
title={Graph-based 3d multi-person pose estimation using multi-view images},
author={Wu, Size and Jin, Sheng and Liu, Wentao and Bai, Lei and Qian, Chen and Liu, Dong and Ouyang, Wanli},
booktitle={Proceedings of the IEEE/CVF international conference on computer vision},
pages={11148--11157},
year={2021}
}
"""
def normalize(mx):
"""Row-normalize sparse matrix"""
rowsum = np.array(mx.sum(1))
r_inv = np.power(rowsum, -1).flatten()
r_inv[np.isinf(r_inv)] = 0.
r_mat_inv = sp.diags(r_inv)
mx = r_mat_inv.dot(mx)
return mx
def sparse_mx_to_torch_sparse_tensor(sparse_mx):
"""Convert a scipy sparse matrix to a torch sparse tensor."""
sparse_mx = sparse_mx.tocoo().astype(np.float32)
indices = torch.from_numpy(np.vstack((sparse_mx.row, sparse_mx.col)).astype(np.int64))
values = torch.from_numpy(sparse_mx.data)
shape = torch.Size(sparse_mx.shape)
return torch.sparse.FloatTensor(indices, values, shape)
def adj_mx_from_edges(num_pts, edges, sparse=True):
edges = np.array(edges, dtype=np.int32)
data, i, j = np.ones(edges.shape[0]), edges[:, 0], edges[:, 1]
adj_mx = sp.coo_matrix((data, (i, j)), shape=(num_pts, num_pts), dtype=np.float32)
adj_mx = adj_mx + adj_mx.T.multiply(adj_mx.T > adj_mx) - adj_mx.multiply(adj_mx.T > adj_mx)
adj_mx = normalize(adj_mx + sp.eye(adj_mx.shape[0]))
if sparse:
adj_mx = sparse_mx_to_torch_sparse_tensor(adj_mx)
else:
adj_mx = torch.tensor(adj_mx.todense(), dtype=torch.float)
return adj_mx
class ChebConv(nn.Module):
"""
The ChebNet convolution operation.
:param in_c: int, number of input channels.
:param out_c: int, number of output channels.
:param K: int, the order of Chebyshev Polynomial.
"""
def __init__(self, in_c, out_c, K, bias=True, normalize=True):
super(ChebConv, self).__init__()
self.normalize = normalize
self.weight = nn.Parameter(torch.Tensor(K + 1, 1, in_c, out_c)) # [K+1, 1, in_c, out_c]
init.xavier_normal_(self.weight)
if bias:
self.bias = nn.Parameter(torch.Tensor(1, 1, out_c))
init.zeros_(self.bias)
else:
self.register_parameter("bias", None)
self.K = K + 1
def forward(self, inputs: torch.Tensor, graph) -> torch.Tensor:
"""
:param inputs: the input data, [B, N, C],
:param graph: the graph structure, [N, N]
:return: convolution result, [B, N, D]
"""
L = ChebConv.get_laplacian(graph, self.normalize).to('cuda') # [N, N]
mul_L = self.cheb_polynomial(L).unsqueeze(1) # [K, 1, N, N]
result = torch.matmul(mul_L, inputs) # [K, B, N, C]
result = torch.matmul(result, self.weight) # [K, B, N, D]
result = torch.sum(result, dim=0) + self.bias # [B, N, D]
return result
def cheb_polynomial(self, laplacian):
"""
Compute the Chebyshev Polynomial, according to the graph laplacian.
:param laplacian: the graph laplacian
:return: the multi order Chebyshev laplacian
"""
N = laplacian.size(0) # [N, N]
multi_order_laplacian = torch.zeros([self.K, N, N], device=laplacian.device, dtype=torch.float) # [K, N, N]
multi_order_laplacian[0] = torch.eye(N, device=laplacian.device, dtype=torch.float)
if self.K == 1:
return multi_order_laplacian
else:
multi_order_laplacian[1] = laplacian
if self.K == 2:
return multi_order_laplacian
else:
for k in range(2, self.K):
multi_order_laplacian[k] = 2 * torch.mm(laplacian, multi_order_laplacian[k-1]) - \
multi_order_laplacian[k-2]
return multi_order_laplacian
@staticmethod
def get_laplacian(graph, normalize):
"""
return the laplacian of the graph.
:param graph: the graph structure without self loop, [N, N].
:param normalize: whether to used the normalized laplacian.
:return: graph laplacian.
"""
if normalize:
D = torch.diag(torch.sum(graph, dim=-1) ** (-1 / 2))
L = torch.eye(graph.size(0), device=graph.device, dtype=graph.dtype) - torch.mm(torch.mm(D, graph), D)
else:
D = torch.diag(torch.sum(graph, dim=-1))
L = D - graph
return L
class _GraphConv(nn.Module):
def __init__(self, input_dim, output_dim, p_dropout=None):
super(_GraphConv, self).__init__()
self.gconv = ChebConv(input_dim, output_dim, K=2)
self.relu = nn.ReLU()
if p_dropout is not None:
self.dropout = nn.Dropout(p_dropout)
else:
self.dropout = None
def forward(self, x, adj):
x = self.gconv(x, adj)
if self.dropout is not None:
x = self.dropout(self.relu(x))
else:
pass
return x
class _ResChebGC(nn.Module):
def __init__(self, input_dim, output_dim, hid_dim, n_seq, p_dropout):
super(_ResChebGC, self).__init__()
self.adj = adj_mx_from_edges(num_pts=n_seq, edges=gan_edges, sparse=False)
self.gconv1 = _GraphConv(input_dim, hid_dim, p_dropout)
def forward(self, x):
out = self.gconv1(x, self.adj)
return out
gan_edges = torch.tensor([[0, 1], [1, 2], [2, 3], [3, 4],
[0, 5], [5, 6], [6, 7], [7, 8],
[0, 9], [9, 10], [10, 11], [11, 12],
[0, 13], [13, 14], [14, 15], [15, 16],
[0, 17], [17, 18], [18, 19], [19, 20]], dtype=torch.long)
body_edges = torch.tensor([[0, 1], [1, 2], [2, 3],
[0, 4], [4, 5], [5, 6],
[0, 7], [7, 8], [8, 9],
[8, 10], [10, 11], [11, 12],
[8, 13], [13, 14], [14, 15]], dtype=torch.long)