def __init__(self, *, input_channels, output_channels, dim, **kwargs):
super().__init__()
self.lin = nn.Sequential(
nn.Linear(input_channels, dim),
nn.ReLU(inplace=True),
nn.Linear(dim, dim),
)
self.unpool = FSPool(dim, 20, relaxed=True)
self.conv = nn.Sequential(
nn.Conv1d(dim, dim, 1),
nn.ReLU(inplace=True),
nn.Conv1d(dim, output_channels, 1),
)
def __init__(self, input_channels, output_channels, dim):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(2 * input_channels, dim, 1),
nn.ReLU(),
nn.Conv2d(dim, output_channels, 1),
)
self.lin = nn.Linear(dim, output_channels)
self.pool = FSPool(output_channels, 20, relaxed=False)
def __init__(self, input_channels, output_channels, dim):
super().__init__()
self.conv = nn.Sequential(
nn.Conv1d(input_channels, dim, 1),
nn.ReLU(),
nn.Conv1d(dim, dim, 1),
nn.ReLU(),
nn.Conv1d(dim, output_channels - 1, 1),
)
self.pool = FSPool(output_channels - 1, 20, relaxed=False)
def __init__(self, dim, point_dim = 128):
super().__init__()
self.enc = nn.Sequential(
nn.Linear(input_channels * set_size, dim)
)
self.lin = nn.Sequential(
nn.Linear(dim, dim, 1),
nn.ReLU(inplace=True),
nn.Linear(dim, output_channels, 1),
)
self.pool = FSPool(dim, 20, relaxed=kwargs.get('relaxed', True))
class FSDecoder(nn.Module):
def __init__(self, *, input_channels, output_channels, dim, **kwargs):
super().__init__()
self.lin = nn.Sequential(
nn.Linear(input_channels, dim),
nn.ReLU(inplace=True),
nn.Linear(dim, dim),
)
self.unpool = FSPool(dim, 20, relaxed=True)
self.conv = nn.Sequential(
nn.Conv1d(dim, dim, 1),
nn.ReLU(inplace=True),
nn.Conv1d(dim, output_channels, 1),
)
def forward(self, x, perm, n_points, *args):
x = self.lin(x)
x, mask = self.unpool.forward_transpose(x, perm, n=n_points)
x = self.conv(x) * mask[:, :1, :]
return x