Pytorch3d: CUDA error when joining meshes as batch

[wboerdijk]

Description

Joining meshes that have textures and are on the (same) gpu with join_meshes_as_batch() results in a CUDA error. The error happens because the image tensor is converted to numpy without calling .cpu() first here:

https://github.com/facebookresearch/pytorch3d/blob/686c8666d31d932ed42d3cd7319f249fc75e89a9/pytorch3d/structures/textures.py#L45

Minimal working example (pytorch3d version = 0.2.0)

import torch
from pytorch3d.structures import Meshes, Textures, join_meshes_as_batch

# create some vertices / faces
verts = torch.randn(1, 10, 3)
faces = torch.randint(0, 10, size=(1, 10, 3))

# create some texture
verts_uvs = torch.randn(1, 10, 2)
faces_uvs = faces
texture_image = torch.randn(1, 100, 100, 3)
texture = Textures(verts_uvs=verts_uvs, faces_uvs=faces_uvs, maps=texture_image)

# create meshes on cpu / gpu
mesh_cpu = Meshes(verts=verts, faces=faces, textures=texture)
mesh_cuda = mesh_cpu.clone().cuda()

# works
joined_mesh_cpu = join_meshes_as_batch([mesh_cpu, mesh_cpu, mesh_cpu])

# CUDA error
joined_mesh_cuda = join_meshes_as_batch([mesh_cuda, mesh_cuda, mesh_cuda])

[nikhilaravi]

Hi @wboerdijk thanks for reporting this. This requires changes to torchvision transforms in order to support GPU tensors for the resize operations. For now you could call join_meshes_as_batch with the meshes on CPU and then move the batch to GPU afterwards e.g.

joined_mesh_cpu = join_meshes_as_batch([mesh_cpu, mesh_cpu, mesh_cpu])
joined_mesh_cuda = joined_mesh_cpu.to("cuda:0")

[nikhilaravi]

@wboerdijk let me know if this suggestion resolved the issue for now. We will add a todo to update this function to support GPU tensors for the texture maps.

[wboerdijk]

Thanks for the quick reply @nikhilaravi! I'm fine with manually moving the mesh, just wanted to let you know about this behavior. You are right about the transforming operations on the gpu of course, that would be quite nice, but I guess we have to wait a bit for that... Issue can be closed from my side!

[nikhilaravi]

@bottler is working on a fix so we will keep this issue open until we have the change landed :)

[bottler]

The fix has landed.