File size: 4,289 Bytes
7a4df11 |
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 |
import torch
import torch.nn as nn
import torch_redstone as rst
from einops import rearrange
from .pointnet_util import PointNetSetAbstraction
class PreNorm(nn.Module):
def __init__(self, dim, fn):
super().__init__()
self.norm = nn.LayerNorm(dim)
self.fn = fn
def forward(self, x, *extra_args, **kwargs):
return self.fn(self.norm(x), *extra_args, **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 = 8, dim_head = 64, dropout = 0., rel_pe = False):
super().__init__()
inner_dim = dim_head * heads
project_out = not (heads == 1 and dim_head == dim)
self.heads = heads
self.scale = dim_head ** -0.5
self.attend = nn.Softmax(dim = -1)
self.dropout = nn.Dropout(dropout)
self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
self.to_out = nn.Sequential(
nn.Linear(inner_dim, dim),
nn.Dropout(dropout)
) if project_out else nn.Identity()
self.rel_pe = rel_pe
if rel_pe:
self.pe = nn.Sequential(nn.Conv2d(3, 64, 1), nn.ReLU(), nn.Conv2d(64, 1, 1))
def forward(self, x, centroid_delta):
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 = self.heads), qkv)
pe = self.pe(centroid_delta) if self.rel_pe else 0
dots = (torch.matmul(q, k.transpose(-1, -2)) + pe) * self.scale
attn = self.attend(dots)
attn = self.dropout(attn)
out = torch.matmul(attn, v)
out = rearrange(out, 'b h n d -> b n (h d)')
return self.to_out(out)
class Transformer(nn.Module):
def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., rel_pe = False):
super().__init__()
self.layers = nn.ModuleList([])
for _ in range(depth):
self.layers.append(nn.ModuleList([
PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout, rel_pe = rel_pe)),
PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
]))
def forward(self, x, centroid_delta):
for attn, ff in self.layers:
x = attn(x, centroid_delta) + x
x = ff(x) + x
return x
class PointPatchTransformer(nn.Module):
def __init__(self, dim, depth, heads, mlp_dim, sa_dim, patches, prad, nsamp, in_dim=3, dim_head=64, rel_pe=False, patch_dropout=0) -> None:
super().__init__()
self.patches = patches
self.patch_dropout = patch_dropout
self.sa = PointNetSetAbstraction(npoint=patches, radius=prad, nsample=nsamp, in_channel=in_dim + 3, mlp=[64, 64, sa_dim], group_all=False)
self.lift = nn.Sequential(nn.Conv1d(sa_dim + 3, dim, 1), rst.Lambda(lambda x: torch.permute(x, [0, 2, 1])), nn.LayerNorm([dim]))
self.cls_token = nn.Parameter(torch.randn(dim))
self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, 0.0, rel_pe)
def forward(self, features):
self.sa.npoint = self.patches
if self.training:
self.sa.npoint -= self.patch_dropout
# print("input", features.shape)
centroids, feature = self.sa(features[:, :3], features)
# print("f", feature.shape, 'c', centroids.shape)
x = self.lift(torch.cat([centroids, feature], dim=1))
x = rst.supercat([self.cls_token, x], dim=-2)
centroids = rst.supercat([centroids.new_zeros(1), centroids], dim=-1)
centroid_delta = centroids.unsqueeze(-1) - centroids.unsqueeze(-2)
x = self.transformer(x, centroid_delta)
return x[:, 0]
class Projected(nn.Module):
def __init__(self, ppat, proj) -> None:
super().__init__()
self.ppat = ppat
self.proj = proj
def forward(self, features: torch.Tensor):
return self.proj(self.ppat(features))
|