Running out of memory, how to split up computation

[blooop]

I am running code similar to this IK example:
https://github.com/NVlabs/curobo/blob/main/examples/ik_example.py

When I query a small number of points <2000 curobo can solve the ik in a single call. I am looking to increase the number of ik samples I investigate but run out of memory. What would your suggestion be to work around memory limitations?

I have looked into the batch parameter but it doesn't seem to be the solution. I have also written some manual batching code that looks at the factors of the ik query list size on the cpu and creates a set of smaller tensors that fit in gpu memory, and recombines the results again on the cpu. This feels overly complicated and dirty.

Is there a way to do this batching natively? My searching on pytorch docs have not come up with anything.

Thanks

[balakumar-s]

I usually create a result buffer for the full batch size and then run a for loop over sub-batch-size, calling IK and filling the result in the result buffer. An example for a collision function:

curobo/src/curobo/graph/graph_base.py

Line 276 in c09d949

for i in range(math.ceil(x_samples.shape[0] / self.max_cg_buffer)):

Not sure if this what you are asking.

[blooop]

I think that answers my question.

This is the solution I have, although I still don't really like it:

def ik(self, frame_list: List[np.ndarray], pos_only: bool = False) -> IKResult:
        """Take a numpy list of 4x4 homogenious matrices and return an IKResult with the solution for each frame"""
        from more_itertools import chunked

        results = []
        for it, chunk in enumerate(chunked(frame_list, 200)):
            if pos_only:
                pose = self.pos(chunk)
            else:
                pose = self.pose_tensor(chunk)
            print(f"solving chunk {it}")
            results.append(self.ik_solver.solve_batch(pose))
        if len(results) == 1:
            return results[0]
        else:
            #combine all the individual IKResults into a combined IKResult
            return IKResult(
                torch.tensor(np.empty((1, 1))), #bug in state.py stopping this gettign js tensor
                torch.tensor(np.empty((1, 1))), #need to look into how to combine poses
                torch.concat([res.solution for res in results]),
                torch.tensor(np.empty((1, 1))), #errors getting seed
                torch.concat([res.success for res in results]),
                torch.concat([res.position_error for res in results]),
                torch.concat([res.rotation_error for res in results]),
                torch.concat([res.error for res in results]),
                sum(res.solve_time for res in results),
            )

I ran into a bug when trying this code but have done a pr for a fix:
#128