| |
| |
| |
| |
| |
|
|
| 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 |
|
|
