From 5953f4a847edcaf8fc52511366ebff6020ce8b09 Mon Sep 17 00:00:00 2001 From: avirupju Date: Tue, 21 Jul 2026 14:35:53 +0000 Subject: [PATCH] Fix PyTorch GPU device-placement bugs so GNM model, landmarks, pruning, and from_gnm run correctly on CUDA. --- gnm/shape/gnm_common.py | 27 ++++++++++++++++++++++++--- gnm/shape/gnm_pytorch.py | 11 +++++++++++ gnm/shape/gnm_pytorch_test.py | 1 + gnm/shape/gnm_xnp.py | 15 +++++++++++++-- 4 files changed, 49 insertions(+), 5 deletions(-) diff --git a/gnm/shape/gnm_common.py b/gnm/shape/gnm_common.py index dd5a3109..bb8437f4 100644 --- a/gnm/shape/gnm_common.py +++ b/gnm/shape/gnm_common.py @@ -22,6 +22,22 @@ _EPSILON = 1e-8 +def _constant_like( + values: Any, + reference_array: enpt.FloatArray['...'], + xnp: Any, + dtype: Any, +) -> enpt.FloatArray['...']: + """Creates a constant array on the same device as `reference_array`. + + For PyTorch this places the constant on `reference_array.device`; for other + backends it falls back to a plain `asarray`. + """ + if enp.lazy.is_torch(reference_array): + return xnp.asarray(values, dtype=dtype, device=reference_array.device) + return xnp.asarray(values, dtype=dtype) + + def take( array: enpt.FloatArray['...'], indices: enpt.IntArray['...'], @@ -48,6 +64,8 @@ def take( xnp = enp.lazy.get_xnp(array) if enp.lazy.is_torch_xnp(xnp): axis = axis % array.ndim + # index_select requires indices to live on the same device as `array`. + indices = xnp.as_tensor(indices, device=array.device) indices_shape = indices.shape if len(indices_shape) > 1: indices_flat = indices.reshape(-1) @@ -184,7 +202,10 @@ def axis_angle_to_rotation_matrix( norm_squared = xnp.sum(xnp.square(axis_angle), axis=-1, keepdims=True) angle = xnp.sqrt( - xnp.maximum(norm_squared, xnp.asarray(epsilon, dtype=norm_squared.dtype)) + xnp.maximum( + norm_squared, + _constant_like(epsilon, axis_angle, xnp, norm_squared.dtype), + ) ) axis = axis_angle / angle @@ -275,8 +296,8 @@ def joint_transforms_world( joints, 2, (num_joints, 1, 4), dtype=joints.dtype ) # Set the last element to 1.0. - bottom_row = bottom_row + xnp.asarray( - [0.0, 0.0, 0.0, 1.0], dtype=joints.dtype + bottom_row = bottom_row + _constant_like( + [0.0, 0.0, 0.0, 1.0], joints, xnp, joints.dtype ) local_transforms = xnp.concatenate( diff --git a/gnm/shape/gnm_pytorch.py b/gnm/shape/gnm_pytorch.py index 76cde071..66d35526 100644 --- a/gnm/shape/gnm_pytorch.py +++ b/gnm/shape/gnm_pytorch.py @@ -131,6 +131,17 @@ def _from_model_data( """Creates a PyTorch GNM instance from model data.""" return cls._from_model_data_with_xnp(model_data, xnp=torch) + @classmethod + def from_gnm(cls, gnm: gnm_xnp.GNM) -> GNM: + """Creates a PyTorch GNM from another GNM, preserving its device.""" + new_gnm = super().from_gnm(gnm) + # `from_gnm` reconstructs the model from a NumPy data dict, so the new model + # is created on CPU. If the source model lives on another device, move the + # reconstructed model there to match. + if isinstance(gnm, torch.nn.Module): + new_gnm = new_gnm.to(gnm.template_vertex_positions.device) + return new_gnm + def compute_vertex_normals( self, vertices: torch.Tensor, diff --git a/gnm/shape/gnm_pytorch_test.py b/gnm/shape/gnm_pytorch_test.py index 59d6e435..cf1dc13e 100644 --- a/gnm/shape/gnm_pytorch_test.py +++ b/gnm/shape/gnm_pytorch_test.py @@ -279,6 +279,7 @@ def test_prune_vertices(self, version: str, variant: Any): gnm_pytorch.GNMMajorVersion(version.removeprefix('v')), gnm_pytorch.GNMVariant(variant_str), ) + gnm_pruned.to(self.device) keep_vertices = gnm_np.quads[0] gnm_pruned.prune_vertices(keep_vertices) diff --git a/gnm/shape/gnm_xnp.py b/gnm/shape/gnm_xnp.py index 396384e9..bfa2bcf3 100644 --- a/gnm/shape/gnm_xnp.py +++ b/gnm/shape/gnm_xnp.py @@ -379,6 +379,10 @@ def vertices_and_landmarks( translation=translation, ) weights = enp.compat.astype(weights, vertices.dtype) + # The landmark tensors are cached lazily and may live on a different device + # than the (possibly moved) model; align them with the output vertices. + if enp.lazy.is_torch(vertices): + weights = weights.to(vertices.device) face_vertices = gnm_common.take(vertices, indices, axis=-2, xnp=self.xnp) landmarks = self.xnp.sum(face_vertices * weights[..., None], axis=-2) @@ -634,6 +638,10 @@ def prune_vertices(self, keep_vertices: enpt.IntArray['V_pruned']) -> None: xnp = self.xnp num_vertices = self.num_vertices keep_vertices = xnp.asarray(keep_vertices, dtype=xnp.int32) + # Keep the index tensor on the same device as the model buffers so that all + # downstream gathers stay device-consistent. + if enp.lazy.is_torch(self.template_vertex_positions): + keep_vertices = keep_vertices.to(self.template_vertex_positions.device) self.template_vertex_positions = gnm_common.take( self.template_vertex_positions, keep_vertices, axis=0, xnp=xnp @@ -729,8 +737,11 @@ def _scatter_indices(keep_vertices, num_vertices, xnp) -> Any: jnp.arange(len(keep_vertices), dtype=jnp.int32) ) elif enp.lazy.is_torch_xnp(xnp): - mapper = xnp.full((num_vertices,), -1, dtype=xnp.int32) - mapper[keep_vertices] = xnp.arange(len(keep_vertices), dtype=xnp.int32) + device = keep_vertices.device + mapper = xnp.full((num_vertices,), -1, dtype=xnp.int32, device=device) + mapper[keep_vertices] = xnp.arange( + len(keep_vertices), dtype=xnp.int32, device=device + ) return mapper else: mapper = np.full((num_vertices,), -1, dtype=np.int32)