def test_grid_sample_backend_equality_2d_batched(self):
grid = math.random_normal(mybatch=10, y=10, x=7)
coords = math.random_uniform(mybatch=10, x=3, y=2) * (12, 9)
grid_ = math.tensor(grid.native('mybatch,x,y'), 'mybatch,x,y')
coords_ = coords.vector.flip()
for extrap in (extrapolation.ZERO, extrapolation.ONE,
extrapolation.BOUNDARY, extrapolation.PERIODIC):
sampled = []
for backend in BACKENDS:
with backend:
grid, coords, grid_, coords_ = math.tensors(
grid, coords, grid_, coords_)
sampled.append(math.grid_sample(grid, coords, extrap))
sampled.append(math.grid_sample(grid_, coords_, extrap))
math.assert_close(*sampled, abs_tolerance=1e-5)
def test_grid_sample(self):
for backend in BACKENDS:
with backend:
grid = math.sum(math.meshgrid(x=[1, 2, 3], y=[0, 3]), 'vector') # 1 2 3 | 4 5 6
coords = math.tensor([(0, 0), (0.5, 0), (0, 0.5), (-2, -1)], instance('list'), channel('vector'))
interp = math.grid_sample(grid, coords, extrapolation.ZERO)
math.assert_close(interp, [1, 1.5, 2.5, 0], msg=backend.name)
def test_grid_sample(self):
for backend in (math.NUMPY_BACKEND, tf.TF_BACKEND,
torch.TORCH_BACKEND):
with backend:
grid = math.sum(math.meshgrid(x=[1, 2, 3], y=[0, 3]),
'vector') # 1 2 3 | 4 5 6
coords = math.tensor([(0, 0), (0.5, 0), (0, 0.5), (-2, -1)],
names=('list', 'vector'))
interp = math.grid_sample(grid, coords, extrapolation.ZERO)
math.assert_close(interp, [1, 1.5, 2.5, 0])
def sample_at(self, points, reduce_channels=()) -> Tensor:
local_points = self.box.global_to_local(points) * self.resolution - 0.5
if len(reduce_channels) == 0:
return math.grid_sample(self.values, local_points,
self.extrapolation)
else:
assert self.shape.channel.sizes == points.shape.get_size(
reduce_channels)
if len(reduce_channels) > 1:
raise NotImplementedError(
f"{len(reduce_channels)} > 1. Only 1 reduced channel allowed."
)
channels = []
for i, channel in enumerate(self.values.vector.unstack()):
channels.append(
math.grid_sample(channel, local_points[{
reduce_channels[0]:
i
}], self.extrapolation))
return math.channel_stack(channels, 'vector')
def test_grid_sample_gradient_1d(self):
for backend in BACKENDS:
if backend.supports(Backend.gradients):
with backend:
grid = math.tensor([0., 1, 2, 3], spatial('x'))
coords = math.tensor([0.5, 1.5], instance('points'))
with math.record_gradients(grid, coords):
sampled = math.grid_sample(grid, coords, extrapolation.ZERO)
loss = math.mean(math.l2_loss(sampled)) / 2
grad_grid, grad_coords = math.gradients(loss, grid, coords)
math.assert_close(grad_grid, math.tensor([0.125, 0.5, 0.375, 0], spatial('x')), msg=backend)
math.assert_close(grad_coords, math.tensor([0.25, 0.75], instance('points')), msg=backend)
def test_grid_sample_gradient_2d(self):
grads_grid = []
grads_coords = []
for backend in BACKENDS:
if backend.supports(Backend.gradients):
with backend:
grid = math.tensor([[1., 2, 3], [1, 2, 3]], spatial('x,y'))
coords = math.tensor([(0.5, 0.5), (1, 1.1), (-0.8, -0.5)], instance('points'), channel('vector'))
with math.record_gradients(grid, coords):
sampled = math.grid_sample(grid, coords, extrapolation.ZERO)
loss = math.sum(sampled) / 3
grad_grid, grad_coords = math.gradients(loss, grid, coords)
grads_grid.append(grad_grid)
grads_coords.append(grad_coords)
math.assert_close(*grads_grid)
math.assert_close(*grads_coords)
def _sample(self, geometry: Geometry) -> Tensor:
if geometry == self.bounds:
return math.mean(self._values, self._resolution)
if isinstance(geometry, GeometryStack):
sampled = [self.sample(g) for g in geometry.geometries]
return math.stack(sampled, geometry.stack_dim)
if isinstance(geometry, GridCell):
if self.elements == geometry:
return self.values
elif math.close(self.dx, geometry.size):
fast_resampled = self._shift_resample(geometry.resolution,
geometry.bounds)
if fast_resampled is not NotImplemented:
return fast_resampled
points = geometry.center
local_points = self.box.global_to_local(points) * self.resolution - 0.5
return math.grid_sample(self.values, local_points, self.extrapolation)
def test_grid_sample_gradient_2d(self):
grads_grid = []
grads_coords = []
for backend in [tf.TF_BACKEND, torch.TORCH_BACKEND]:
with backend:
grid = math.tensor([[1., 2, 3], [1, 2, 3]], 'x,y')
coords = math.tensor([(0.5, 0.5), (1, 1.1), (-0.8, -0.5)],
'points,vector')
with math.record_gradients(grid, coords):
sampled = math.grid_sample(grid, coords,
extrapolation.ZERO)
loss = math.sum(sampled)
grad_grid, grad_coords = math.gradients(loss, grid, coords)
grads_grid.append(grad_grid)
grads_coords.append(grad_coords)
math.assert_close(*grads_grid)
math.assert_close(*grads_coords)
def test_grid_sample_gradient_1d(self):
grads_grid = []
grads_coords = []
for backend in BACKENDS:
if backend.supports(Backend.gradients):
print(backend)
with backend:
grid = math.tensor([0., 1, 2, 3], 'x')
coords = math.tensor([0.5, 1.5], 'points')
with math.record_gradients(grid, coords):
sampled = math.grid_sample(grid, coords,
extrapolation.ZERO)
loss = math.l2_loss(sampled)
grad_grid, grad_coords = math.gradients(
loss, grid, coords)
grads_grid.append(grad_grid)
grads_coords.append(grad_coords)
math.assert_close(*grads_grid, math.tensor([0.125, 0.5, 0.375, 0],
'x'))
math.assert_close(*grads_coords, math.tensor([0.25, 0.75], 'points'))
def test_grid_sample_1d(self):
grid = math.tensor([0, 1, 2, 3], names='x')
coords = math.tensor([[0], [1], [0.5]], names='x,vector')
sampled = math.grid_sample(grid, coords, None)
math.print(sampled)
math.assert_close(sampled, [0, 1, 0.5])
def test_grid_sample_1d(self):
grid = math.tensor([0, 1, 2, 3], spatial('x'))
coords = math.tensor([[0], [1], [0.5]], spatial('x'), channel('vector'))
sampled = math.grid_sample(grid, coords, None)
math.print(sampled)
math.assert_close(sampled, [0, 1, 0.5])
