NVIDIA / warp

A Python framework for high performance GPU simulation and graphics
https://nvidia.github.io/warp/
Other
4.21k stars 239 forks source link

About gradients clipping in Tape #180

Closed Leon-LXA closed 4 months ago

Leon-LXA commented 8 months ago

Hi, I am currently doing the gradients clipping in Tape.py. At first I tried to use warp.clamp(), but it can only be used for a scalar. Gradients in tape are in the form of array so I wanna use a for loop to clamp each element of the array. But it shows that 'RuntimeError: Item indexing is not supported on wp.array objects'. I wanna know if there's any other method to do it, thx :D

Best regards, Lei

c0d1f1ed commented 8 months ago

Hi Lei, you can call wp.clamp() for each array element in a @wp.kernel function. Typically you can add that operation to the end of the kernel that computes the array values.

Leon-LXA commented 8 months ago

Hi Lei, you can call wp.clamp() for each array element in a @wp.kernel function. Typically you can add that operation to the end of the kernel that computes the array values.

I am trying to modify tape.py. However in get_adjoint function, it seems like vector a would have different types and it is unable for me to use clamp for each element in the array. Could you please show me a simple modification example in tape.py? Thanks a lot!

mmacklin commented 6 months ago

Hi Leon, if I understand correctly you want to clip just the gradients (not the forward mode of the computation, correct?) - you might be able to achieve this through a custom gradient function: https://nvidia.github.io/warp/modules/runtime.html#example-1-custom-grad-function

Otherwise you could register custom function to the wp.Tape using https://nvidia.github.io/warp/modules/runtime.html#warp.Tape.record_func, this will be called during the backward pass and I think you should be able to launch a new kernel to perform the clamping there.

Hope that helps.

Cheers, Miles