Spaces:
Sleeping
Sleeping
| Interoperability | |
| ================ | |
| Warp can interop with other Python-based frameworks such as NumPy through standard interface protocols. | |
| NumPy | |
| ----- | |
| Warp arrays may be converted to a NumPy array through the ``warp.array.numpy()`` method. When the Warp array lives on | |
| the ``cpu`` device this will return a zero-copy view onto the underlying Warp allocation. If the array lives on a | |
| ``cuda`` device then it will first be copied back to a temporary buffer and copied to NumPy. | |
| Warp CPU arrays also implement the ``__array_interface__`` protocol and so can be used to construct NumPy arrays | |
| directly:: | |
| w = wp.array([1.0, 2.0, 3.0], dtype=float, device="cpu") | |
| a = np.array(w) | |
| print(a) | |
| > [1. 2. 3.] | |
| .. _pytorch-interop: | |
| PyTorch | |
| ------- | |
| Warp provides helper functions to convert arrays to/from PyTorch. Please see the ``warp.torch`` module for more details. Example usage is shown below:: | |
| import warp.torch | |
| w = wp.array([1.0, 2.0, 3.0], dtype=float, device="cpu") | |
| # convert to Torch tensor | |
| t = warp.to_torch(w) | |
| # convert from Torch tensor | |
| w = warp.from_torch(t) | |
| These helper functions allow the conversion of Warp arrays to/from PyTorch tensors without copying the underlying data. | |
| At the same time, if available, gradient arrays and tensors are converted to/from PyTorch autograd tensors, allowing the use of Warp arrays | |
| in PyTorch autograd computations. | |
| Example: Optimization using ``warp.from_torch()`` | |
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ | |
| An example usage of minimizing a loss function over an array of 2D points written in Warp via PyTorch's Adam optimizer using ``warp.from_torch`` is as follows:: | |
| import warp as wp | |
| import torch | |
| wp.init() | |
| @wp.kernel() | |
| def loss(xs: wp.array(dtype=float, ndim=2), l: wp.array(dtype=float)): | |
| tid = wp.tid() | |
| wp.atomic_add(l, 0, xs[tid, 0] ** 2.0 + xs[tid, 1] ** 2.0) | |
| # indicate requires_grad so that Warp can accumulate gradients in the grad buffers | |
| xs = torch.randn(100, 2, requires_grad=True) | |
| l = torch.zeros(1, requires_grad=True) | |
| opt = torch.optim.Adam([xs], lr=0.1) | |
| wp_xs = wp.from_torch(xs) | |
| wp_l = wp.from_torch(l) | |
| tape = wp.Tape() | |
| with tape: | |
| # record the loss function kernel launch on the tape | |
| wp.launch(loss, dim=len(xs), inputs=[wp_xs], outputs=[wp_l], device=wp_xs.device) | |
| for i in range(500): | |
| tape.zero() | |
| tape.backward(loss=wp_l) # compute gradients | |
| # now xs.grad will be populated with the gradients computed by Warp | |
| opt.step() # update xs (and thereby wp_xs) | |
| # these lines are only needed for evaluating the loss | |
| # (the optimization just needs the gradient, not the loss value) | |
| wp_l.zero_() | |
| wp.launch(loss, dim=len(xs), inputs=[wp_xs], outputs=[wp_l], device=wp_xs.device) | |
| print(f"{i}\tloss: {l.item()}") | |
| Example: Optimization using ``warp.to_torch`` | |
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ | |
| Less code is needed when we declare the optimization variables directly in Warp and use ``warp.to_torch`` to convert them to PyTorch tensors. | |
| Here, we revisit the same example from above where now only a single conversion to a torch tensor is needed to supply Adam with the optimization variables:: | |
| import warp as wp | |
| import numpy as np | |
| import torch | |
| wp.init() | |
| @wp.kernel() | |
| def loss(xs: wp.array(dtype=float, ndim=2), l: wp.array(dtype=float)): | |
| tid = wp.tid() | |
| wp.atomic_add(l, 0, xs[tid, 0] ** 2.0 + xs[tid, 1] ** 2.0) | |
| # initialize the optimization variables in Warp | |
| xs = wp.array(np.random.randn(100, 2), dtype=wp.float32, requires_grad=True) | |
| l = wp.zeros(1, dtype=wp.float32, requires_grad=True) | |
| # just a single wp.to_torch call is needed, Adam optimizes using the Warp array gradients | |
| opt = torch.optim.Adam([wp.to_torch(xs)], lr=0.1) | |
| tape = wp.Tape() | |
| with tape: | |
| wp.launch(loss, dim=len(xs), inputs=[xs], outputs=[l], device=xs.device) | |
| for i in range(500): | |
| tape.zero() | |
| tape.backward(loss=l) | |
| opt.step() | |
| l.zero_() | |
| wp.launch(loss, dim=len(xs), inputs=[xs], outputs=[l], device=xs.device) | |
| print(f"{i}\tloss: {l.numpy()[0]}") | |
| .. automodule:: warp.torch | |
| :members: | |
| :undoc-members: | |
| CuPy/Numba | |
| ---------- | |
| Warp GPU arrays support the ``__cuda_array_interface__`` protocol for sharing data with other Python GPU frameworks. | |
| Currently this is one-directional, so that Warp arrays can be used as input to any framework that also supports the | |
| ``__cuda_array_interface__`` protocol, but not the other way around. | |
| .. _jax-interop: | |
| JAX | |
| --- | |
| Interoperability with JAX arrays is supported through the following methods. | |
| Internally these use the DLPack protocol to exchange data in a zero-copy way with JAX. | |
| .. automodule:: warp.jax | |
| :members: | |
| :undoc-members: | |
| DLPack | |
| ------ | |
| .. automodule:: warp.dlpack | |
| :members: | |
| :undoc-members: | |