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])