def SetSDFs(self, sdfs, sdf_sizes):
cell_sizes = [
ec.check_conditions(x, "sdf_sizes[%d]" % i, "%s > 0",
"isinstance(%s, numbers.Real)")
for i, x in enumerate(sdf_sizes)
]
_sdfs = [
ec.check_conditions(sdf, "sdfs[%d]" % i,
"isinstance(%s, torch.Tensor)",
"len(%s.size()) == " + str(self.ndim))
for i, sdf in enumerate(sdfs)
]
_sdf_shapes = ec.list2tensor(
[list(x.size()) + [
cell_sizes[i],
] for i, x in enumerate(_sdfs)])
_sdfs = [x.contiguous().view(-1) for x in _sdfs]
_sdf_offsets = ec.list2tensor([
0,
] + np.cumsum([x.size()[0] for x in _sdfs])[:-1].tolist())
_sdfs = torch.cat(_sdfs)
self.sdfs.resize_(_sdfs.size())
self.sdfs[...] = _sdfs
self.sdf_shapes.resize_(_sdf_shapes.size())
self.sdf_shapes[...] = _sdf_shapes
self.sdf_offsets.resize_(_sdf_offsets.size())
self.sdf_offsets[...] = _sdf_offsets
def __init__(self,
in_channels,
out_channels,
ndim,
kernel_size,
dilation,
radius,
dis_norm=False,
kernel_fn='default',
with_params=True):
""" Initialize a Smooth Particle Convolution layer.
Arguments:
-in_channels: The number of features for each input particle.
-out_channels: The number of features to output for each particle.
-ndim: The dimensionality of the particle's coordinate space.
-kernel_size: (int or tuple) The shape of the kernel that is place around each
particle. The kernel is centered on the particle, so the size
must be odd.
-dilation: (float or tuple) The spacing between each cell of the kernel.
-radius: The radius to use when computing the neighbors for each query point.
-kernel_fn: The kernel function to use in the SPH equation. Refer to kernels.py
for a list and explanation of all available functions.
-dis_norm: If true, will divide by the particle-to-particle distance in the
SPH equation.
-with_params: If true, the parameters weight and bias are registered with
PyTorch as parameters. Otherwise they are registered as buffers,
meaning they won't be optimized when doing backprop.
"""
super(ConvSP, self).__init__()
self.nchannels = ec.check_conditions(
in_channels, "in_channels", "%s > 0",
"isinstance(%s, numbers.Integral)")
self.nkernels = ec.check_conditions(
out_channels, "out_channels", "%s > 0",
"isinstance(%s, numbers.Integral)")
self.ndim = ec.check_conditions(
ndim, "ndim", "%s > 0",
"%s < " + str(_ext.spn_max_cartesian_dim()),
"isinstance(%s, numbers.Integral)")
self._kernel_size = ec.make_list(kernel_size, ndim, "kernel_size",
"%s >= 0",
"%s %% 2 == 1 # Must be odd",
"isinstance(%s, numbers.Integral)")
self._dilation = ec.make_list(dilation, ndim, "dilation", "%s >= 0",
"isinstance(%s, numbers.Real)")
self.radius = ec.check_conditions(radius, "radius", "%s >= 0",
"isinstance(%s, numbers.Real)")
self.kernel_fn = ec.check_conditions(kernel_fn, "kernel_fn",
"%s in " + str(KERNEL_NAMES))
self.kernel_fn = KERNEL_NAMES.index(self.kernel_fn)
self.dis_norm = (1 if dis_norm else 0)
self.ncells = np.prod(self._kernel_size)
if with_params:
self.register_parameter(
"weight",
torch.nn.Parameter(
torch.Tensor(self.nkernels, self.nchannels, self.ncells)))
self.register_parameter(
"bias", torch.nn.Parameter(torch.Tensor(self.nkernels)))
else:
self.register_buffer(
"weight",
torch.autograd.Variable(
torch.Tensor(self.nkernels, self.nchannels, self.ncells)))
self.register_buffer(
"bias", torch.autograd.Variable(torch.Tensor(self.nkernels)))
self.register_buffer("kernel_size", ec.list2tensor(self._kernel_size))
self.register_buffer("dilation", ec.list2tensor(self._dilation))
self.nshared_device_mem = -1
self.device_id = -1
def __init__(self,
sdfs,
sdf_sizes,
out_channels,
ndim,
kernel_size,
dilation,
max_distance,
with_params=True,
compute_pose_grads=False):
""" Initialize a SDF Convolution layer.
Arguments:
-sdfs: List of SDFs. Each SDF must be a ndim-dimensional float tensor, with
the value of the SDF evaluated at the center of each cell in the grid.
-sdf_sizes: List of the size of one side of a grid cell in each SDf in sdfs.
-out_channels: The number of features to output for each query location.
-ndim: The dimensionality of the coordinate space.
-kernel_size: (int or tuple) The shape of the kernel that is place around each
query location. The kernel is centered on the location, so the
size must be odd.
-dilation: (float or tuple) The spacing between each cell of the kernel.
-max_distance: A cap on the maximum SDF value, i.e., the SDF value at any
point p is min(min_i SDF_i(p), max_distance).
-with_params: If true, the parameters weight and bias are registered with
PyTorch as parameters. Otherwise they are registered as buffers,
meaning they won't be optimized when doing backprop.
-compute_pose_grads: If False, will not compute grads wrt to the sdfposes during
backpropagation. This is done to increase speed when
gradients are not neede for the sdf poses.
"""
super(ConvSDF, self).__init__()
self.nkernels = ec.check_conditions(
out_channels, "out_channels", "%s > 0",
"isinstance(%s, numbers.Integral)")
self.ndim = ec.check_conditions(
ndim, "ndim", "%s > 0",
"%s < " + str(_ext.spn_max_cartesian_dim()),
"%s in [1, 2, 3] # Only 1-, 2-, and 3-D are suported",
"isinstance(%s, numbers.Integral)")
self.max_distance = ec.check_conditions(
max_distance, "max_distance", "%s >= 0",
"isinstance(%s, numbers.Real)")
self._kernel_size = ec.make_list(kernel_size, ndim, "kernel_size",
"%s >= 0",
"%s %% 2 == 1 # Must be odd",
"isinstance(%s, numbers.Integral)")
self._dilation = ec.make_list(dilation, ndim, "dilation", "%s >= 0",
"isinstance(%s, numbers.Real)")
self.register_buffer("sdfs", torch.zeros(1))
self.register_buffer("sdf_shapes", torch.zeros(1))
self.register_buffer("sdf_offsets", torch.zeros(1))
self.SetSDFs(sdfs, sdf_sizes)
self.ncells = np.prod(self._kernel_size)
if with_params:
self.register_parameter(
"weight",
torch.nn.Parameter(torch.Tensor(self.nkernels, self.ncells)))
self.register_parameter(
"bias", torch.nn.Parameter(torch.Tensor(self.nkernels)))
else:
self.register_buffer(
"weight",
torch.autograd.Variable(
torch.Tensor(self.nkernels, self.ncells)))
self.register_buffer(
"bias", torch.autograd.Variable(torch.Tensor(self.nkernels)))
self.compute_pose_grads = (True if compute_pose_grads else False)
self._kernel_size = ec.list2tensor(self._kernel_size)
self._dilation = ec.list2tensor(self._dilation)
self.register_buffer("kernel_size", self._kernel_size)
self.register_buffer("dilation", self._dilation)