Spaces:
Sleeping
Sleeping
| import torch.nn as nn | |
| import torch | |
| from einops import rearrange | |
| import numpy | |
| def index_points(device, points, idx): | |
| """ | |
| Input: | |
| points: input points data, [B, N, C] | |
| idx: sample index data, [B, S] | |
| Return: | |
| new_points:, indexed points data, [B, S, C] | |
| """ | |
| B = points.shape[0] | |
| view_shape = list(idx.shape) | |
| view_shape[1:] = [1] * (len(view_shape) - 1) | |
| repeat_shape = list(idx.shape) | |
| repeat_shape[0] = 1 | |
| # batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape) | |
| batch_indices = torch.arange(B, dtype=torch.long).cuda().view(view_shape).repeat(repeat_shape) | |
| new_points = points[batch_indices, idx, :] | |
| return new_points | |
| def knn_l2(device, net, k, u): | |
| ''' | |
| Input: | |
| k: int32, number of k in k-nn search | |
| net: (batch_size, npoint, c) float32 array, points | |
| u: int32, block size | |
| Output: | |
| idx: (batch_size, npoint, k) int32 array, indices to input points | |
| ''' | |
| INF = 1e8 | |
| batch_size = net.size(0) | |
| npoint = net.size(1) | |
| n_channel = net.size(2) | |
| square = torch.pow(torch.norm(net, dim=2,keepdim=True),2) | |
| def u_block(batch_size, npoint, u): | |
| block = numpy.zeros([batch_size, npoint, npoint]) | |
| n = npoint // u | |
| for i in range(n): | |
| block[:, (i*u):(i*u+u), (i*u):(i*u+u)] = numpy.ones([batch_size, u, u]) * (-INF) | |
| return block | |
| # minus_distance = 2 * torch.matmul(net, net.transpose(2,1)) - square - square.transpose(2,1) + torch.Tensor(u_block(batch_size, npoint, u)).to(device) | |
| minus_distance = 2 * torch.matmul(net, net.transpose(2,1)) - square - square.transpose(2,1) + torch.Tensor(u_block(batch_size, npoint, u)).cuda() | |
| _, indices = torch.topk(minus_distance, k, largest=True, sorted=False) | |
| return indices | |
| class Residual(nn.Module): | |
| def __init__(self, fn): | |
| super().__init__() | |
| self.fn = fn | |
| def forward(self, x, **kwargs): | |
| return self.fn(x, **kwargs) + x | |
| class PreNorm(nn.Module): | |
| def __init__(self, dim, fn): | |
| super().__init__() | |
| self.norm = nn.LayerNorm(dim) | |
| self.fn = fn | |
| def forward(self, x, **kwargs): | |
| return self.fn(self.norm(x), **kwargs) | |
| class FeedForward(nn.Module): | |
| def __init__(self, dim, hidden_dim, dropout = 0.): | |
| super().__init__() | |
| self.net = nn.Sequential( | |
| nn.Linear(dim, hidden_dim), | |
| nn.GELU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(hidden_dim, dim), | |
| nn.Dropout(dropout) | |
| ) | |
| def forward(self, x): | |
| return self.net(x) | |
| class Attention(nn.Module): | |
| def __init__(self, dim, heads = 4, dropout = 0.): | |
| super().__init__() | |
| self.heads = heads | |
| self.scale = dim ** -0.5 | |
| self.to_qkv = nn.Linear(dim, dim * 3, bias = False) | |
| self.to_out = nn.Sequential( | |
| nn.Linear(dim, dim), | |
| nn.Dropout(dropout) | |
| ) | |
| def forward(self, x, mask = None): | |
| b, n, _, h = *x.shape, self.heads | |
| qkv = self.to_qkv(x).chunk(3, dim = -1) | |
| q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv) | |
| dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale | |
| if mask is not None: | |
| mask = F.pad(mask.flatten(1), (1, 0), value = True) | |
| assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions' | |
| mask = mask[:, None, :] * mask[:, :, None] | |
| dots.masked_fill_(~mask, float('-inf')) | |
| del mask | |
| attn = dots.softmax(dim=-1) | |
| out = torch.einsum('bhij,bhjd->bhid', attn, v) | |
| out = rearrange(out, 'b h n d -> b n (h d)') | |
| out = self.to_out(out) | |
| return out | |
| class Local_Attention(nn.Module): | |
| def __init__(self, dim, heads = 4,knn=4, dropout = 0.): | |
| super().__init__() | |
| self.heads = heads | |
| self.scale = dim ** -0.5 | |
| #self.to_qkv = nn.Linear(dim, dim * 3, bias = False) | |
| self.q=nn.Linear(dim,dim,bias=False) | |
| self.k=nn.Linear(dim,dim,bias=False) | |
| self.v=nn.Linear(dim,dim,bias=False) | |
| self.to_out = nn.Sequential( | |
| nn.Linear(dim, dim), | |
| nn.Dropout(dropout) | |
| ) | |
| self.knn=knn | |
| def forward(self, x, mask = None): | |
| b, n, _, h = *x.shape, self.heads | |
| point=x*1 | |
| X=x*1 | |
| idx = knn_l2(0, point.permute(0,2,1), 4, 1) | |
| feat=idx | |
| new_point = index_points(0, point.permute(0,2,1),idx) | |
| group_point = new_point.permute(0, 3, 2, 1) | |
| _1,_2,_3,_4=group_point.shape | |
| q=self.q(X.reshape(_1*_2,1,_4)) | |
| k=self.k(torch.cat([group_point.reshape(_1*_2,self.knn,_4),X.reshape(_1*_2,1,_4)],dim=1)) | |
| v=self.v(torch.cat([group_point.reshape(_1*_2,self.knn,_4),X.reshape(_1*_2,1,_4)],dim=1)) | |
| q, k, v = rearrange(q, 'b n (h d) -> b h n d', h = h),rearrange(k, 'b n (h d) -> b h n d', h = h),rearrange(v, 'b n (h d) -> b h n d', h = h) | |
| attn_map=q@k.permute(0,1,3,2)*self.scale | |
| attn_map=attn_map.softmax(dim=-1) | |
| out=attn_map@v | |
| out=out.view(b,out.shape[0]//b,out.shape[1],out.shape[3]).permute(0,2,1,3) | |
| out = rearrange(out, 'b h n d -> b n (h d)') | |
| out = self.to_out(out) | |
| return out | |
| class Transformer(nn.Module): | |
| def __init__(self, dim, depth, heads, mlp_dim, group=5, dropout=0.): | |
| super().__init__() | |
| self.layers = nn.ModuleList([]) | |
| for _ in range(depth): | |
| self.layers.append(nn.ModuleList([ | |
| Residual(PreNorm(dim, Attention(dim, heads = heads, dropout = dropout))), | |
| Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))) | |
| ])) | |
| self.group=group | |
| def forward(self, x, mask = None): | |
| bs_gp,dim,wid,hei=x.shape[0],x.shape[1],x.shape[2],x.shape[3] | |
| bs=bs_gp//self.group | |
| gp=self.group | |
| x=x.reshape(bs,gp,dim,wid,hei) | |
| x=x.permute(0,1,3,4,2).reshape(bs,gp*wid*hei,dim) | |
| for attn, ff in self.layers: | |
| x = attn(x, mask = mask) | |
| x = ff(x) | |
| x=x.reshape(bs,gp,wid,hei,dim).permute(0,1,4,2,3).reshape(bs_gp,dim,wid,hei) | |
| return x | |
| class Transformer__local(nn.Module): | |
| def __init__(self, dim, depth, heads, mlp_dim,knn_k=4, group=5, dropout=0.): | |
| super().__init__() | |
| self.layers = nn.ModuleList([]) | |
| for _ in range(depth): | |
| self.layers.append(nn.ModuleList([ | |
| Residual(PreNorm(dim, Local_Attention(dim, heads = heads,knn=knn_k, dropout = dropout))), | |
| Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))) | |
| ])) | |
| self.group=group | |
| def forward(self, x, mask = None): | |
| bs_gp,dim,wid,hei=x.shape[0],x.shape[1],x.shape[2],x.shape[3] | |
| bs=bs_gp//self.group | |
| gp=self.group | |
| x=x.reshape(bs,gp,dim,wid,hei) | |
| x=x.permute(0,1,3,4,2).reshape(bs,gp*wid*hei,dim) | |
| for attn, ff in self.layers: | |
| x = attn(x, mask = mask) | |
| x = ff(x) | |
| x=x.reshape(bs,gp,wid,hei,dim).permute(0,1,4,2,3).reshape(bs_gp,dim,wid,hei) | |
| return x |