| | |
| | |
| | |
| | |
| | |
| |
|
| | import torch |
| | import torch.nn as nn |
| | import torch.sparse as sp |
| |
|
| |
|
| | |
| | class LocalAffine(nn.Module): |
| |
|
| | def __init__(self, num_points, batch_size=1, edges=None): |
| | ''' |
| | specify the number of points, the number of points should be constant across the batch |
| | and the edges torch.Longtensor() with shape N * 2 |
| | the local affine operator supports batch operation |
| | batch size must be constant |
| | add additional pooling on top of w matrix |
| | ''' |
| | super(LocalAffine, self).__init__() |
| | self.A = nn.Parameter( |
| | torch.eye(3).unsqueeze(0).unsqueeze(0).repeat( |
| | batch_size, num_points, 1, 1)) |
| | self.b = nn.Parameter( |
| | torch.zeros(3).unsqueeze(0).unsqueeze(0).unsqueeze(3).repeat( |
| | batch_size, num_points, 1, 1)) |
| | self.edges = edges |
| | self.num_points = num_points |
| |
|
| | def stiffness(self): |
| | ''' |
| | calculate the stiffness of local affine transformation |
| | f norm get infinity gradient when w is zero matrix, |
| | ''' |
| | if self.edges is None: |
| | raise Exception("edges cannot be none when calculate stiff") |
| | idx1 = self.edges[:, 0] |
| | idx2 = self.edges[:, 1] |
| | affine_weight = torch.cat((self.A, self.b), dim=3) |
| | w1 = torch.index_select(affine_weight, dim=1, index=idx1) |
| | w2 = torch.index_select(affine_weight, dim=1, index=idx2) |
| | w_diff = (w1 - w2)**2 |
| | w_rigid = (torch.linalg.det(self.A) - 1.0)**2 |
| | return w_diff, w_rigid |
| |
|
| | def forward(self, x, return_stiff=False): |
| | ''' |
| | x should have shape of B * N * 3 |
| | ''' |
| | x = x.unsqueeze(3) |
| | out_x = torch.matmul(self.A, x) |
| | out_x = out_x + self.b |
| | out_x.squeeze_(3) |
| | if return_stiff: |
| | stiffness, rigid = self.stiffness() |
| | return out_x, stiffness, rigid |
| | else: |
| | return out_x |
| |
|
