CUDA out of memory

[ZisongXu]

Thanks for the nice work. When I tried mustard0, and it works fine. But when I try to run other objects such as driller (example in your link) and cheezit (collected by myself), it had this error:

Traceback (most recent call last):
  File "run_demo.py", line 53, in <module>
    pose = est.register(K=reader.K, rgb=color, depth=depth, ob_mask=mask, iteration=args.est_refine_iter)
  File "/root/FoundationPose/estimater.py", line 219, in register
    scores, vis = self.scorer.predict(mesh=self.mesh, rgb=rgb, depth=depth, K=K, ob_in_cams=poses.data.cpu().numpy(), normal_map=normal_map, mesh_tensors=self.mesh_tensors, glctx=self.glctx, mesh_diameter=self.diameter, get_vis=self.debug>=2)
  File "/opt/conda/envs/my/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/FoundationPose/learning/training/predict_score.py", line 185, in predict
    pose_data = make_crop_data_batch(self.cfg.input_resize, ob_in_cams, mesh, rgb, depth, K, crop_ratio=self.cfg['crop_ratio'], glctx=glctx, mesh_tensors=mesh_tensors, dataset=self.dataset, cfg=self.cfg, mesh_diameter=mesh_diameter)
  File "/opt/conda/envs/my/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/FoundationPose/learning/training/predict_score.py", line 110, in make_crop_data_batch
    pose_data = dataset.transform_batch(pose_data, H_ori=H, W_ori=W, bound=1)
  File "/root/FoundationPose/learning/datasets/h5_dataset.py", line 178, in transform_batch
    batch = self.transform_depth_to_xyzmap(batch, H_ori, W_ori, bound=bound)
  File "/root/FoundationPose/learning/datasets/h5_dataset.py", line 159, in transform_depth_to_xyzmap
    depthBs_ori = kornia.geometry.transform.warp_perspective(batch.depthBs.cuda().expand(bs,-1,-1,-1), crop_to_oris, dsize=(H_ori, W_ori), mode='nearest', align_corners=False)
  File "/opt/conda/envs/my/lib/python3.8/site-packages/kornia/geometry/transform/imgwarp.py", line 124, in warp_perspective
    grid = transform_points(src_norm_trans_dst_norm[:, None, None], grid)
  File "/opt/conda/envs/my/lib/python3.8/site-packages/kornia/geometry/linalg.py", line 192, in transform_points
    points_0_h = torch.bmm(points_1_h, trans_01.permute(0, 2, 1))
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.60 GiB (GPU 0; 9.77 GiB total capacity; 5.23 GiB already allocated; 2.47 GiB free; 5.35 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

I want to ask is that really because CUDA out of memory, or because I did something wrong when I change objects.

By the way, if I want to run some experiments use my own data, what should I do? I only changed two lines in run_demo.py script:

parser.add_argument('--mesh_file', type=str, default=f'{code_dir}/demo_data/mustard0/mesh/textured_simple.obj')
parser.add_argument('--test_scene_dir', type=str, default=f'{code_dir}/demo_data/mustard0')

Best Regards Zisong XU

[abhishekmonogram]

@ZisongXu It happened with me too when I was working with higher resolution images. The initial pose estimation algorithm takes a lot of GPU but only for a couple of seconds and then it drops. So try resizing the rgb, depth and masks to a smaller resolution and then run. Make sure to change cam_K.txt by the same downscale factor while you resize them. That should help you run it on your machine. I am not sure if there is any other way around it. @wenbowen123 can suggest something maybe

[wenbowen123]

@ZisongXu how much GPU memory do you have?

[ZisongXu]

@wenbowen123 My GPU is 3080, memory is 10G I think

[wenbowen123]

https://github.com/NVlabs/FoundationPose/blob/7b484cf03436b9bb6038768e6edd4bc06d759513/run_demo.py#L44 You can set shorter_side to something smaller, e.g. 480

[ZisongXu]

Nice it works! Thanks a lot!!