可微性#

默认情况下，Warp 为每个内核定义生成一个前向版本和一个后向（伴随）版本。内核的后向版本可用于计算损失函数的梯度，这些梯度可以反向传播到 PyTorch 等机器学习框架。

参与计算链且需要梯度的数组必须使用 requires_grad=True 创建，例如

a = wp.zeros(1024, dtype=wp.vec3, device="cuda", requires_grad=True)

然后可以使用 wp.Tape 类来记录内核启动并重放它们，以计算标量损失函数相对于内核输入的梯度。

tape = wp.Tape()

# forward pass
with tape:
    wp.launch(kernel=compute1, inputs=[a, b], device="cuda")
    wp.launch(kernel=compute2, inputs=[c, d], device="cuda")
    wp.launch(kernel=loss, inputs=[d, l], device="cuda")

# reverse pass
tape.backward(l)

后向传播完成后，输入相对于的梯度可以从 array.grad 属性获取。

# gradient of loss with respect to input a
print(a.grad)

注意，梯度会累积在参与的缓冲区上，因此如果希望在多次后向传播中重用相同的缓冲区，应首先使用 Tape.zero() 将梯度清零。

class warp.Tape[源]#

在 Tape 范围内记录内核启动以启用自动微分。在磁带上记录操作后，可以通过 Tape.backward() 计算梯度。

示例

tape = wp.Tape()

# forward pass
with tape:
    wp.launch(kernel=compute1, inputs=[a, b], device="cuda")
    wp.launch(kernel=compute2, inputs=[c, d], device="cuda")
    wp.launch(kernel=loss, inputs=[d, l], device="cuda")

# reverse pass
tape.backward(l)

梯度可以通过 tape.gradients 字典访问，例如

print(tape.gradients[a])
__init__()[源]#
backward(loss=None, grads=None)[源]#

评估磁带上记录的操作的后向传播。可以提供一个单元素数组 loss 或一个数组字典 grads，以分配用于反向模式自动微分的传入梯度。

参数：

  • loss (wp.array) – 一个单元素数组，包含要计算其梯度的损失函数值。

  • grads (dict) – 一个数组字典，将 Warp 数组映射到其传入的梯度。

record_launch(
kernel,
dim,
max_blocks,
inputs,
outputs,
device,
block_dim=0,
metadata=None,
)[源]#
record_func(backward, arrays)[源]#

记录一个自定义函数，该函数仅在后向传播中执行。

参数：

  • backward (Callable) – 一个可调用的 Python 对象（可以是任何函数），将在后向传播中执行。

  • arrays (list) – 后向函数使用的数组列表。磁带会跟踪这些数组，以便能够在 Tape.zero() 中将它们的梯度清零。

record_scope_begin(scope_name, metadata=None)[源]#

在磁带上开始一个范围，以便将操作分组。范围仅用于可视化函数。

record_scope_end(remove_scope_if_empty=True)[源]#

结束磁带上的范围。

参数：

remove_scope_if_empty (bool) – 如果为 True，则如果范围内没有记录内核启动，该范围将被移除。

get_adjoint(a)[源]#
reset()[源]#

清除磁带上记录的所有操作并清零所有梯度。

zero()[源]#

清零磁带上记录的所有梯度。

visualize(
filename=None,
simplify_graph=True,
hide_readonly_arrays=False,
array_labels=None,
choose_longest_node_name=True,
ignore_graph_scopes=False,
track_inputs=None,
track_outputs=None,
track_input_names=None,
track_output_names=None,
graph_direction='LR',
)[源]#

将磁带上记录的操作可视化为 GraphViz 图。

示例

import warp as wp

tape = wp.Tape()
with tape:
    # record Warp kernel launches here
    wp.launch(...)

dot_code = tape.visualize("tape.dot")

此函数创建一个 GraphViz dot 文件，可以使用 GraphViz 命令行工具将其渲染为图像，例如通过

dot -Tpng tape.dot -o tape.png
参数：

  • filename (str) – 保存可视化文件的文件名（可选）。

  • simplify_graph (bool) – 如果为 True，则通过检测重复的内核启动序列并在子图中进行总结来简化图。

  • hide_readonly_arrays (bool) – 如果为 True，则隐藏未被任何内核启动修改的数组。

  • array_labels (Dict[wp.array, str]) – 将数组映射到自定义标签的字典。

  • choose_longest_node_name (bool) – 如果为 True，自动名称解析将尝试为计算图中的每个数组找到最长的名称。

  • ignore_graph_scopes (bool) – 如果为 True，则在可视化图时忽略磁带上记录的范围。

  • track_inputs (List[wp.array]) – 在图中作为输入跟踪的数组列表，以确保无论 hide_readonly_arrays 设置如何，它们都会显示。

  • track_outputs (List[wp.array]) – 在图中作为输出跟踪的数组列表，以确保它们保持可见。

  • track_input_names (List[str]) – 在图中跟踪的输入数组的自定义名称列表（与 track_inputs 结合使用）。

  • track_output_names (List[str]) – 在图中跟踪的输出数组的自定义名称列表（与 track_outputs 结合使用）。

  • graph_direction (str) – 图布局的方向（默认：“LR”）。

返回：

表示图的 dot 代码。

返回类型：

str

数组覆盖#

为了正确计算梯度，自动微分框架必须存储计算的中间结果用于反向传播。覆盖之前计算的结果可能导致梯度计算不正确。因此，PyTorch 和 JAX 等框架会为每个操作输出隐式分配新内存。在 Warp 中，用户显式管理内存，因此在使用 tape.backward() 等功能时应注意避免覆盖之前的结果。

考虑 PyTorch 中的以下梯度计算

import torch

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2 + 3 * x + 1
y.backward(torch.ones_like(x))

print(x.grad)

或者，等效地，在 JAX 中

import jax
import jax.numpy as jnp

def func(x):
    return x ** 2 + 3 * x + 1

x = jnp.array([1.0, 2.0, 3.0])
grad_func = jax.vmap(jax.grad(func))
x_grad = grad_func(x)

print(x_grad)

这两个框架都只要求用户显式分配张量 x：y 和 x_grad 是通过赋值隐式分配的。在 Warp 中，我们可以这样写

import warp as wp

@wp.kernel
def kernel_func(x: wp.array(dtype=float), y: wp.array(dtype=float)):
    tid = wp.tid()
    y[tid] = x[tid] ** 2.0 + 3.0 * x[tid] + 1.0

x = wp.array([1.0, 2.0, 3.0], dtype=float, requires_grad=True)
y = wp.zeros_like(x)

tape = wp.Tape()
with tape:
    wp.launch(kernel_func, x.shape, inputs=[x], outputs=[y])
tape.backward(grads={y: wp.ones_like(x)})

print(x.grad)

[5. 7. 9.]

x 和 y 在开始时显式分配。注意，我们可以写 wp.launch(kernel_func, x.shape, inputs=[x,y])，但有时使用 wp.launch() 中的 inputs 和 outputs 参数来跟踪哪些数组被读取/写入很有用（实际上，在可视化计算图时这样做是必要的）。

import warp as wp
wp.config.enable_backward = False

@wp.kernel
def kernel_func(x: wp.array(dtype=float)):
    tid = wp.tid()
    x[tid] = x[tid] ** 2.0 + 3.0 * x[tid] + 1.0

x = wp.array([1.0, 2.0, 3.0], dtype=float)

wp.launch(kernel_func, x.shape, inputs=[x])

print(x)

[ 5. 11. 19.]

如果不需要 x 的梯度和先验值，我们可以写成 这样只需要四分之一的内存分配，但这会使梯度跟踪无效。

尤其对于大型计算图，很难辨别数组是否被覆盖。在这种情况下，将 wp.config.verify_autograd_array_access=True 设置为 True 会很有帮助，它将自动检测数组覆盖。在此处阅读更多。

注意

尽管原地操作，例如 x[tid] += 1.0 和 wp.atomic_add()，从技术上讲是覆盖操作，但 Warp 图特别适应了这些情况下的伴随累加。在此处阅读更多。

复制是可微分的#

wp.copy()、wp.clone() 和 array.assign() 是可微分函数，可以参与磁带上记录的计算图。考虑以下示例及其 PyTorch 等效写法（用于比较）

wp.copy():

@wp.kernel
def double(x: wp.array(dtype=float), y: wp.array(dtype=float)):
    tid = wp.tid()
    y[tid] = x[tid] * 2.0

x = wp.array(np.arange(3), dtype=float, requires_grad=True)
y = wp.zeros_like(x)
z = wp.zeros_like(x)

tape = wp.Tape()
with tape:
    wp.launch(double, dim=3, inputs=[x, y])
    wp.copy(z, y)

tape.backward(grads={z: wp.ones_like(x)})

print(x.grad)

[2. 2. 2.]

等效地，在 PyTorch 中

x = torch.tensor(np.arange(3), dtype=torch.float32, requires_grad=True)
y = x * 2
z = torch.zeros_like(y).copy_(y)

z.sum().backward()

print(x.grad)
# tensor([2., 2., 2.])

wp.clone():

x = wp.array(np.arange(3), dtype=float, requires_grad=True)
y = wp.zeros_like(x)

tape = wp.Tape()
with tape:
    wp.launch(double, dim=3, inputs=[x, y])
    z = wp.clone(y, requires_grad=True)

tape.backward(grads={z: wp.ones_like(x)})

print(x.grad)

[2. 2. 2.]

在 PyTorch 中

x = torch.tensor(np.arange(3), dtype=torch.float32, requires_grad=True)
y = x * 2
z = torch.clone(y)

z.sum().backward()
print(x.grad)
# tensor([2., 2., 2.])

注意

在 PyTorch 中，可以通过调用 x.clone().detach() 克隆张量 x 并将其从当前计算图中分离。在 Warp 中，等效写法是 wp.clone(x, requires_grad=False)。

array.assign():

x = wp.array(np.arange(3), dtype=float, requires_grad=True)
y = wp.zeros_like(x)
z = wp.zeros_like(y)

tape = wp.Tape()
with tape:
    wp.launch(double, dim=3, inputs=[x], outputs=[y])
    z.assign(y)

tape.backward(grads={z: wp.ones_like(x)})

print(x.grad)

[2. 2. 2.]

注意

array.assign() 等效于 wp.copy()，外加一个步骤，如果源数组不是 Warp 数组，则将其包装在 Warp 数组中。

雅可比矩阵#

为了计算多值函数 \(f: \mathbb{R}^n \to \mathbb{R}^m\) 的雅可比矩阵 \(J\in\mathbb{R}^{m\times n}\)，我们可以通过计算雅可比矩阵向量积 \(J^\top \mathbf{e}\) 来并行评估雅可比矩阵的整行。向量 \(\mathbf{e}\in\mathbb{R}^m\) 选择输出缓冲区中要对其进行微分的索引。在 Warp 中，我们不将标量损失缓冲区传递给 tape.backward() 方法，而是传递一个字典 grads，将函数输出数组映射到具有相同类型的选择向量 \(\mathbf{e}\)。

# compute the Jacobian for a function of single output
jacobian = np.empty((output_dim, input_dim), dtype=np.float32)

# record computation
tape = wp.Tape()
with tape:
    output_buffer = launch_kernels_to_be_differentiated(input_buffer)

# compute each row of the Jacobian
for output_index in range(output_dim):

    # select which row of the Jacobian we want to compute
    select_index = np.zeros(output_dim)
    select_index[output_index] = 1.0
    e = wp.array(select_index, dtype=wp.float32)

    # pass input gradients to the output buffer to apply selection
    tape.backward(grads={output_buffer: e})
    q_grad_i = tape.gradients[input_buffer]
    jacobian[output_index, :] = q_grad_i.numpy()

    # zero gradient arrays for next row
    tape.zero()

当我们独立并行运行模拟时，对应于整个系统动力学的雅可比矩阵是一个分块对角矩阵。在这种情况下，我们可以通过选择一个对所有环境副本的输出索引都有效的选择向量，来并行计算所有环境的雅可比矩阵。例如，为了获得所有环境雅可比矩阵的第一行，\(\mathbf{e}=[\begin{smallmatrix}1 & 0 & 0 & \dots & 1 & 0 & 0 & \dots\end{smallmatrix}]^\top\)，计算第二行则使用 \(\mathbf{e}=[\begin{smallmatrix}0 & 1 & 0 & \dots & 0 & 1 & 0 & \dots\end{smallmatrix}]^\top\) 等。

# compute the Jacobian for a function over multiple environments in parallel
jacobians = np.empty((num_envs, output_dim, input_dim), dtype=np.float32)

# record computation
tape = wp.Tape()
with tape:
    output_buffer = launch_kernels_to_be_differentiated(input_buffer)

# compute each row of the Jacobian
for output_index in range(output_dim):

    # select which row of the Jacobian we want to compute
    select_index = np.zeros(output_dim)
    select_index[output_index] = 1.0

    # assemble selection vector for all environments (can be precomputed)
    e = wp.array(np.tile(select_index, num_envs), dtype=wp.float32)
    tape.backward(grads={output_buffer: e})
    q_grad_i = tape.gradients[input_buffer]
    jacobians[:, output_index, :] = q_grad_i.numpy().reshape(num_envs, input_dim)

    tape.zero()

自定义梯度函数#

Warp 支持为用户定义的 Warp 函数定义自定义梯度函数。这允许用户定义应替代自动生成的导数的代码。

为了微分包含对函数 \(g(x)\) 的嵌套调用的函数 \(h(x) = f(g(x))\)，在 \(h(x)\) 的自动微分中评估链式法则

\[h^\prime(x) = f^\prime({\color{green}{\underset{\textrm{replay}}{g(x)}}}) {\color{blue}{\underset{\textrm{grad}}{g^\prime(x)}}}\]

这意味着一个函数要与自动微分引擎兼容，需要提供其前向版本 \(\color{green}{g(x)}\) 的实现，我们称之为“重放”（replay）函数（默认匹配原始函数定义），以及其导数 \(\color{blue}{g^\prime(x)}\)，称为“梯度”（grad）。

重放和梯度实现都可以由用户自定义。它们的定义如下

自定义函数 myfunc 的重放和梯度版本#

前向函数

 
@wp.func
def myfunc(in1: InType1, ..., inN: InTypeN) -> OutType1, ..., OutTypeM:
    return out1, ..., outM

自定义重放函数

 
@wp.func_replay(myfunc)
def replay_myfunc(in1: InType1, ..., inN: InTypeN) -> OutType1, ..., OutTypeM:
    # Custom forward computations to be executed in the backward pass of a
    # function calling `myfunc` go here
    # Ensure the output variables match the original forward definition
    return out1, ..., outM

自定义梯度函数

 
@wp.func_grad(myfunc)
def adj_myfunc(in1: InType1, ..., inN: InTypeN, adj_out1: OutType1, ..., adj_outM: OutTypeM):
    # Custom adjoint code goes here
    # Update the partial derivatives for the inputs as follows:
    wp.adjoint[in1] += ...
    ...
    wp.adjoint[inN] += ...

注意

目前无法为具有泛型参数（例如 Any 或 wp.array(dtype=Any)）的函数定义自定义重放或梯度函数。使用泛型参数的重放或梯度函数本身也尚不受支持。

示例 1：自定义梯度函数#

在下文中，我们定义一个计算数字平方根的 Warp 函数 safe_sqrt

@wp.func
def safe_sqrt(x: float):
    return wp.sqrt(x)

为了评估此函数，我们定义了一个将 safe_sqrt 应用于输入值数组的内核

@wp.kernel
def run_safe_sqrt(xs: wp.array(dtype=float), output: wp.array(dtype=float)):
    i = wp.tid()
    output[i] = safe_sqrt(xs[i])

对值数组 [1.0, 2.0, 0.0] 调用内核会产生预期的输出，并且除了零输入外，梯度是有限的

xs = wp.array([1.0, 2.0, 0.0], dtype=wp.float32, requires_grad=True)
ys = wp.zeros_like(xs)

with wp.Tape() as tape:
    wp.launch(run_safe_sqrt, dim=len(xs), inputs=[xs], outputs=[ys])
tape.backward(grads={ys: wp.array(np.ones(len(xs)), dtype=wp.float32)})

print("ys     ", ys)
print("xs.grad", xs.grad)
tape.zero()

ys      [1.        1.4142135 0.       ]
xs.grad [0.5        0.35355338        inf]

通常需要在计算图中捕获非有限梯度，因为它们可能导致整个梯度计算变为非有限。为此，我们可以定义一个自定义梯度函数，通过 @wp.func_grad(safe_sqrt) 装饰器来装饰自定义梯度代码，从而替换为 safe_sqrt 自动生成的伴随函数。

@wp.func_grad(safe_sqrt)
def adj_safe_sqrt(x: float, adj_ret: float):
    if x > 0.0:
        wp.adjoint[x] += 1.0 / (2.0 * wp.sqrt(x)) * adj_ret

注意

自定义梯度函数的签名包含前向函数的输入参数以及前向函数输出的伴随变量。要访问和修改输入参数的偏导数，我们使用 wp.adjoint 字典。此字典的键是前向函数的输入参数，值是前向函数输出相对于输入参数的偏导数。

现在，上面代码的输出是

ys      [1.        1.4142135 0.       ]
xs.grad [0.5        0.35355338 0.        ]

示例 2：自定义重放函数#

在下文中，我们通过 wp.atomic_add() 在每个线程中递增一个数组索引，并在递增后的索引处计算输入数组的平方根

@wp.kernel
def test_add(counter: wp.array(dtype=int), input: wp.array(dtype=float), output: wp.array(dtype=float)):
    idx = wp.atomic_add(counter, 0, 1)
    output[idx] = wp.sqrt(input[idx])


dim = 8
use_reversible_increment = False
input = wp.array(np.arange(1, dim + 1), dtype=wp.float32, requires_grad=True)
counter = wp.zeros(1, dtype=wp.int32)
thread_ids = wp.zeros(dim, dtype=wp.int32)
output = wp.zeros(dim, dtype=wp.float32, requires_grad=True)
with wp.Tape() as tape:
    if use_reversible_increment:
        wp.launch(test_add_diff, dim, inputs=[counter, thread_ids, input], outputs=[output])
    else:
        wp.launch(test_add, dim, inputs=[counter, input], outputs=[output])

print("counter:    ", counter.numpy())
print("thread_ids: ", thread_ids.numpy())
print("input:      ", input.numpy())
print("output:     ", output.numpy())

tape.backward(grads={output: wp.array(np.ones(dim), dtype=wp.float32)})
print("input.grad: ", input.grad.numpy())
tape.zero()

上面代码的输出是

counter:     [8]
thread_ids:  [0 0 0 0 0 0 0 0]
input:       [1. 2. 3. 4. 5. 6. 7. 8.]
output:      [1.        1.4142135 1.7320508 2.        2.236068  2.4494898 2.6457512
 2.828427 ]
input.grad:  [4. 0. 0. 0. 0. 0. 0. 0.]

输入的梯度不正确，因为涉及原子操作 wp.atomic_add() 的后向传播不知道哪个线程 ID 对应于哪个输入值。wp.atomic_add() 的伴随函数返回的索引始终为零，因此梯度是输入数组的第一个条目，即 \(\frac{1}{2\sqrt{1}} = 0.5\)，被累加 dim 次（因此 input.grad[0] == 4.0，而所有其他条目为零）。

为了解决这个问题，我们定义了一个新的 Warp 函数 reversible_increment()，它具有一个自定义的重放定义，将线程 ID 存储在一个单独的数组中

@wp.func
def reversible_increment(
    buf: wp.array(dtype=int),
    buf_index: int,
    value: int,
    thread_values: wp.array(dtype=int),
    tid: int
):
    next_index = wp.atomic_add(buf, buf_index, value)
    # store which thread ID corresponds to which index for the backward pass
    thread_values[tid] = next_index
    return next_index


@wp.func_replay(reversible_increment)
def replay_reversible_increment(
    buf: wp.array(dtype=int),
    buf_index: int,
    value: int,
    thread_values: wp.array(dtype=int),
    tid: int
):
    return thread_values[tid]

现在，调用 reversible_increment() 的函数的后向传播中的前向阶段会执行 replay_reversible_increment() 中的自定义重放代码，而不是运行 reversible_increment()。我们在前向传播中首先将数组索引存储到每个线程 ID，现在在后向传播中为每个线程 ID 检索数组索引。这样，后向传播可以以完全相同的线程操作数重现与前向传播相同的加法操作。

警告

自定义重放函数的签名必须与前向函数的签名匹配。

要使用我们的函数，我们编写以下内核

@wp.kernel
def test_add_diff(
    counter: wp.array(dtype=int),
    thread_ids: wp.array(dtype=int),
    input: wp.array(dtype=float),
    output: wp.array(dtype=float)
):
    tid = wp.tid()
    idx = reversible_increment(counter, 0, 1, thread_ids, tid)
    output[idx] = wp.sqrt(input[idx])

通过之前的 main 函数并设置 use_reversible_increment = True 来运行 test_add_diff 内核，我们现在可以计算输入数组的正确梯度

counter:     [8]
thread_ids:  [0 1 2 3 4 5 6 7]
input:       [1. 2. 3. 4. 5. 6. 7. 8.]
output:      [1.        1.4142135 1.7320508 2.        2.236068  2.4494898 2.6457512
 2.828427 ]
input.grad:  [0.5        0.35355338 0.28867513 0.25       0.2236068  0.20412414
 0.18898225 0.17677669]

自定义原生函数#

用户可以使用 @func_native 装饰的函数在 Warp 内核中插入原生 C++/CUDA 代码。这些函数接受原生代码字符串，这些字符串在代码生成后被编译，并在 @wp.kernel 函数中调用。例如

snippet = """
    __shared__ int sum[128];

    sum[tid] = arr[tid];
    __syncthreads();

    for (int stride = 64; stride > 0; stride >>= 1) {
        if (tid < stride) {
            sum[tid] += sum[tid + stride];
        }
        __syncthreads();
    }

    if (tid == 0) {
        out[0] = sum[0];
    }
    """

@wp.func_native(snippet)
def reduce(arr: wp.array(dtype=int), out: wp.array(dtype=int), tid: int): ...


@wp.kernel
def reduce_kernel(arr: wp.array(dtype=int), out: wp.array(dtype=int)):
    tid = wp.tid()
    reduce(arr, out, tid)


N = 128
x = wp.array(np.arange(N, dtype=int), dtype=int)
out = wp.zeros(1, dtype=int)

wp.launch(kernel=reduce_kernel, dim=N, inputs=[x, out])

print(out)

[8128]

注意这里共享内存的使用：Warp 库不暴露共享内存作为功能，但 CUDA 编译器会轻松接受上述代码片段。这意味着 Warp 中未暴露的 CUDA 功能仍然可以在 Warp 脚本中访问。面向 CPU 的 Warp 内核当然无法利用 CUDA 功能，但这个相同的机制也支持纯 C++ 代码片段。

请记住以下几点：代码片段中的线程索引应在 @wp.kernel 中计算并传递给您的代码片段，如上例所示。这意味着您的 @wp.func_native 函数签名应包含您代码片段中使用的变量，以及一个 int 类型的线程索引。函数体本身应使用 ... 进行桩化（代码片段将在编译期间插入）。

如果您希望在磁带上记录原生函数并在之后回卷磁带，则必须在您的代码片段旁边包含一个伴随代码片段作为装饰器的附加输入，如下例所示

snippet = """
out[tid] = a * x[tid] + y[tid];
"""
adj_snippet = """
adj_a += x[tid] * adj_out[tid];
adj_x[tid] += a * adj_out[tid];
adj_y[tid] += adj_out[tid];
"""

@wp.func_native(snippet, adj_snippet)
def saxpy(
    a: wp.float32,
    x: wp.array(dtype=wp.float32),
    y: wp.array(dtype=wp.float32),
    out: wp.array(dtype=wp.float32),
    tid: int,
):
    ...

@wp.kernel
def saxpy_kernel(
    a: wp.float32,
    x: wp.array(dtype=wp.float32),
    y: wp.array(dtype=wp.float32),
    out: wp.array(dtype=wp.float32)
):
    tid = wp.tid()
    saxpy(a, x, y, out, tid)

N = 128
a = 2.0
x = wp.array(np.arange(N, dtype=np.float32), dtype=wp.float32, requires_grad=True)
y = wp.zeros_like(x)
out = wp.array(np.arange(N, dtype=np.float32), dtype=wp.float32)
adj_out = wp.array(np.ones(N, dtype=np.float32), dtype=wp.float32)

tape = wp.Tape()

with tape:
    wp.launch(kernel=saxpy_kernel, dim=N, inputs=[a, x, y], outputs=[out])

tape.backward(grads={out: adj_out})

print(f"x.grad = {x.grad}")
print(f"y.grad = {y.grad}")

x.grad = [2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2.
 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2.
 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2.
 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2.
 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2.
 2. 2. 2. 2. 2. 2. 2. 2.]
y.grad = [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.
 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.
 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.
 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.
 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.
 1. 1. 1. 1. 1. 1. 1. 1.]

您还可以包含一个自定义重放代码片段，作为伴随的一部分执行（有关完整解释，请参阅自定义梯度函数）。考虑以下示例

num_threads = 8
counter = wp.zeros(1, dtype=wp.int32)
thread_values = wp.zeros(num_threads, dtype=wp.int32)
inputs = wp.array(np.arange(num_threads, dtype=np.float32), requires_grad=True)
outputs = wp.zeros_like(inputs)

snippet = """
    int next_index = atomicAdd(counter, 1);
    thread_values[tid] = next_index;
    """
replay_snippet = ""


@wp.func_native(snippet, replay_snippet=replay_snippet)
def reversible_increment(counter: wp.array(dtype=int), thread_values: wp.array(dtype=int), tid: int):
    ...


@wp.kernel
def run_atomic_add(
    input: wp.array(dtype=float),
    counter: wp.array(dtype=int),
    thread_values: wp.array(dtype=int),
    output: wp.array(dtype=float),
):
    tid = wp.tid()
    reversible_increment(counter, thread_values, tid)
    idx = thread_values[tid]
    output[idx] = input[idx] ** 2.0


with wp.Tape() as tape:
    wp.launch(run_atomic_add, dim=num_threads, inputs=[inputs, counter, thread_values], outputs=[outputs])

tape.backward(grads={outputs: wp.ones(num_threads, dtype=wp.float32)})

print(f"inputs.grad = {np.round(inputs.grad.numpy(), 5)}")

inputs.grad = [ 0.  2.  4.  6.  8. 10. 12. 14.]

默认情况下，snippet 会在后向传播中调用，但在本例中，我们定义了一个自定义重放代码片段来代替调用它。replay_snippet 是一个无操作，这正是我们需要的，因为 thread_values 在前向传播中被缓存。如果我们没有定义 replay_snippet，thread_values 在后向传播中会被超过输入数组大小的计数器值覆盖。

原生代码片段也可能包含 return 语句。如果是这种情况，您必须在原生函数定义中指定返回类型，如下例所示

snippet = """
    float sq = x * x;
    return sq;
    """
adj_snippet = """
    adj_x += 2.f * x * adj_ret;
    """


@wp.func_native(snippet, adj_snippet)
def square(x: float) -> float:
    ...


@wp.kernel
def square_kernel(input: wp.array(dtype=Any), output: wp.array(dtype=Any)):
    tid = wp.tid()
    x = input[tid]
    output[tid] = square(x)


N = 5
x = wp.array(np.arange(N, dtype=float), dtype=float, requires_grad=True)
y = wp.zeros_like(x)

with wp.Tape() as tape:
    wp.launch(kernel=square_kernel, dim=N, inputs=[x, y])

tape.backward(grads={y: wp.ones(N, dtype=float)})

print(f"x.grad = {x.grad}")

x.grad = [0. 2. 4. 6. 8.]

调试梯度#

Warp 提供实用函数来评估给定给内核启动的输入/输出参数对的偏雅可比矩阵。jacobian() 使用 Warp 的自动微分引擎计算 Warp 内核或任何调用 Warp 内核并以 Warp 数组作为输入和输出的 Python 函数的雅可比矩阵。jacobian_fd() 使用有限差分计算内核或函数的雅可比矩阵。gradcheck() 比较自动微分引擎和有限差分计算的雅可比矩阵，以衡量梯度的准确性。jacobian_plot() 可视化 jacobian() 和 jacobian_fd() 函数返回的雅可比矩阵。

warp.autograd.gradcheck#

warp.autograd.gradcheck(
function,
dim=None,
inputs=None,
outputs=None,
*,
eps=1e-4,
atol=1e-3,
rtol=1e-2,
raise_exception=True,
input_output_mask=None,
device=None,
max_blocks=0,
block_dim=256,
max_inputs_per_var=-1,
max_outputs_per_var=-1,
plot_relative_error=False,
plot_absolute_error=False,
show_summary=True,
)[source]#

检查 Warp 核函数的自动微分梯度是否与有限差分匹配。

给定自动微分（\(\nabla_\text{AD}\)）梯度和有限差分（\(\nabla_\text{FD}\)）梯度，如果自动微分梯度不包含 NaN 值且满足以下条件，则检查成功：

\[|\nabla_\text{AD} - \nabla_\text{FD}| \leq atol + rtol \cdot |\nabla_\text{FD}|.\]

使用给定的输入和输出以及提供的 dim、max_blocks 和 block_dim 参数启动核函数及其伴随版本（详情请参阅 warp.launch()）。

注意

此函数仅支持输入参数位于输出参数之前的 Warp 核函数。

仅考虑 requires_grad=True 的 Warp 数组用于雅可比计算。

此函数尚不支持 Struct 类型的参数来计算雅可比矩阵。

参数：

  • function (Kernel | Callable) – Warp 核函数，使用 @wp.kernel 装饰器修饰，或者任何涉及 Warp 核函数启动的函数。

  • dim (Tuple[int]) – 启动核函数的线程数，可以是整数或整数元组。仅当函数是 Warp 核函数时才需要。

  • inputs (Sequence) – 输入变量列表。

  • outputs (Sequence) – 输出变量列表。仅当函数是 Warp 核函数时才需要。

  • eps (float) – 有限差分步长。

  • atol (float) – 梯度检查的绝对容差。

  • rtol (float) – 梯度检查的相对容差。

  • raise_exception (bool) – 如果为 True，则在梯度检查失败时引发 ValueError。

  • input_output_mask (List[Tuple[str | int, str | int]]) – 指定要计算雅可比矩阵的输入-输出对的元组列表。输入和输出可以通过它们在核函数输入/输出参数中出现的整数索引来标识，或通过相应的参数名字符串来标识。如果为 None，则计算所有输入-输出对的雅可比矩阵。

  • device (Device | str | None) – 启动设备（可选）。

  • max_blocks (int) – 使用的 CUDA 线程块的最大数量。

  • block_dim (int) – 每个块的线程数。

  • max_inputs_per_var (int) – 评估输入-输出对雅可比矩阵时考虑的最大输入维度数。如果值 <= 0，则评估所有输入维度。

  • max_outputs_per_var (int) – 评估输入-输出对雅可比矩阵时考虑的最大输出维度数。如果值 <= 0，则评估所有输出维度。

  • plot_relative_error (bool) – 如果为 True，则在图中可视化雅可比矩阵的相对误差（需要 matplotlib）。

  • plot_absolute_error (bool) – 如果为 True，则在图中可视化雅可比矩阵的绝对误差（需要 matplotlib）。

  • show_summary (bool) – 如果为 True，则打印梯度检查结果的摘要表。

返回：

如果梯度检查通过，则为 True，否则为 False。

返回类型：

bool

warp.autograd.gradcheck_tape#

warp.autograd.gradcheck_tape(
tape,
*,
eps=1e-4,
atol=1e-3,
rtol=1e-2,
raise_exception=True,
input_output_masks=None,
blacklist_kernels=None,
whitelist_kernels=None,
max_inputs_per_var=-1,
max_outputs_per_var=-1,
plot_relative_error=False,
plot_absolute_error=False,
show_summary=True,
reverse_launches=False,
skip_to_launch_index=0,
)[source]#

检查记录在 Warp Tape 上的核函数的自动微分梯度是否与有限差分匹配。

给定自动微分（\(\nabla_\text{AD}\)）梯度和有限差分（\(\nabla_\text{FD}\)）梯度，如果自动微分梯度不包含 NaN 值且满足以下条件，则检查成功：

\[|\nabla_\text{AD} - \nabla_\text{FD}| \leq atol + rtol \cdot |\nabla_\text{FD}|.\]

注意

仅检查记录在 tape 上的 Warp 核函数，而不检查已记录的任意函数，例如通过 Tape.record_func()。

仅考虑 requires_grad=True 的 Warp 数组用于雅可比计算。

此函数尚不支持 Struct 类型的参数来计算雅可比矩阵。

参数：

  • tape (Tape) – 执行梯度检查的 Warp tape。

  • eps – 有限差分步长。

  • atol – 梯度检查的绝对容差。

  • rtol – 梯度检查的相对容差。

  • raise_exception – 如果为 True，则在梯度检查失败时引发 ValueError。

  • input_output_masks (Dict[str, List[Tuple[str | int, str | int]]]) – Tape 中每个核函数的输入-输出掩码字典，将核函数键映射到输入-输出掩码。输入和输出可以通过它们在核函数输入/输出参数中出现的整数索引来标识，或通过相应的参数名字符串来标识。如果为 None，则计算所有输入-输出对的雅可比矩阵。

  • blacklist_kernels (List[str]) – 要从梯度检查中排除的核函数键列表。

  • whitelist_kernels (List[str]) – 要包含在梯度检查中的核函数键列表。如果非空或非 None，则仅检查此列表中的核函数。

  • max_inputs_per_var – 评估输入-输出对雅可比矩阵时考虑的最大输入维度数。如果值 <= 0，则评估所有输入维度。

  • max_outputs_per_var – 评估输入-输出对雅可比矩阵时考虑的最大输出维度数。如果值 <= 0，则评估所有输出维度。

  • plot_relative_error – 如果为 True，则在图中可视化雅可比矩阵的相对误差（需要 matplotlib）。

  • plot_absolute_error – 如果为 True，则在图中可视化雅可比矩阵的绝对误差（需要 matplotlib）。

  • show_summary (bool) – 如果为 True，则打印梯度检查结果的摘要表。

  • reverse_launches (bool) – 如果为 True，则反转 Tape 上要检查的核函数启动顺序。

  • skip_to_launch_index (int)

返回：

如果 Tape 上所有核函数的梯度检查通过，则为 True，否则为 False。

返回类型：

bool

warp.autograd.jacobian#

warp.autograd.jacobian(
function,
dim=None,
inputs=None,
outputs=None,
input_output_mask=None,
device=None,
max_blocks=0,
block_dim=256,
max_outputs_per_var=-1,
plot_jacobians=False,
metadata=None,
kernel=None,
)[source]#

计算函数或 Warp 核函数对于提供的可微分输入到可微分输出选择的雅可比矩阵。

输入函数可以是 Warp 核函数（例如使用 @wp.kernel 装饰的函数），也可以是接受参数（其中一些必须是 Warp 数组）并返回一个或多个 Warp 数组的普通 Python 函数。

如果 function 是 Warp 核函数，则使用给定的输入和输出以及提供的 dim、max_blocks 和 block_dim 参数启动其伴随核函数（详情请参阅 warp.launch()）。

注意

如果 function 是 Warp 核函数，则输入参数必须在核函数代码定义中位于输出参数之前。

仅考虑 requires_grad=True 的 Warp 数组用于雅可比计算。

此函数尚不支持类型为 Struct 的函数参数。

参数：

  • function (Kernel | Callable) – Warp 核函数，或返回 Warp 数组或 Warp 数组列表的普通 Python 函数。

  • dim (Tuple[int]) – 启动核函数的线程数，可以是整数或整数元组。仅当 function 是 Warp 核函数时才需要。

  • inputs (Sequence) – 输入变量列表。至少有一个参数必须是具有 requires_grad=True 的 Warp 数组。

  • outputs (Sequence) – 输出变量列表。如果函数是返回 Warp 数组或 Warp 数组列表的普通 Python 函数，则为可选。仅当 function 是 Warp 核函数时才需要。

  • input_output_mask (List[Tuple[str | int, str | int]]) – 指定要计算雅可比矩阵的输入-输出对的元组列表。输入和输出可以通过它们在核函数输入/输出参数中出现的整数索引来标识，或通过相应的参数名字符串来标识。如果为 None，则计算所有输入-输出对的雅可比矩阵。

  • device (Device | str | None) – 启动设备（可选）。仅当 function 是 Warp 核函数时使用。

  • max_blocks – 使用的 CUDA 线程块的最大数量。仅当 function 是 Warp 核函数时使用。

  • block_dim – 每个块的线程数。仅当 function 是 Warp 核函数时使用。

  • max_outputs_per_var – 评估输入-输出对雅可比矩阵时考虑的最大输出维度数。如果值 <= 0，则评估所有输出维度。

  • plot_jacobians – 如果为 True，则在图中可视化计算出的雅可比矩阵（需要 matplotlib）。

  • metadata (FunctionMetadata) – 核函数的元数据，包含输入和输出标签、步长和数据类型。如果为 None 或为空，则从核函数或函数推断元数据。

  • kernel (Kernel) – 已弃用的参数。请改用 function 参数。

返回：

一个雅可比矩阵字典，其中键是输入和输出索引的元组，值是雅可比矩阵。

返回类型：

Dict[Tuple[int, int], array]

warp.autograd.jacobian_fd#

warp.autograd.jacobian_fd(
function,
dim=None,
inputs=None,
outputs=None,
input_output_mask=None,
device=None,
max_blocks=0,
block_dim=256,
max_inputs_per_var=-1,
eps=1e-4,
plot_jacobians=False,
metadata=None,
kernel=None,
)[source]#

计算函数或 Warp 核函数对于提供的可微分输入到可微分输出选择的有限差分雅可比矩阵。该方法使用中心差分方案来近似雅可比矩阵。

输入函数可以是 Warp 核函数（例如使用 @wp.kernel 装饰的函数），也可以是接受参数（其中一些必须是 Warp 数组）并返回一个或多个 Warp 数组的普通 Python 函数。

以仅前向模式多次启动函数，并使用给定的输入。如果 function 是 Warp 核函数，则将提供的输入和输出以及其他参数 dim、max_blocks 和 block_dim 提供给核函数启动（详情请参阅 warp.launch()）。

注意

如果 function 是 Warp 核函数，则输入参数必须在核函数代码定义中位于输出参数之前。

仅考虑 requires_grad=True 的 Warp 数组用于雅可比计算。

此函数尚不支持类型为 Struct 的函数参数。

参数：

  • function (Kernel | Callable) – Warp 核函数，或返回 Warp 数组或 Warp 数组列表的普通 Python 函数。

  • dim (Tuple[int]) – 启动核函数的线程数，可以是整数或整数元组。仅当 function 是 Warp 核函数时才需要。

  • inputs (Sequence) – 输入变量列表。至少有一个参数必须是具有 requires_grad=True 的 Warp 数组。

  • outputs (Sequence) – 输出变量列表。如果函数是返回 Warp 数组或 Warp 数组列表的普通 Python 函数，则为可选。仅当 function 是 Warp 核函数时才需要。

  • input_output_mask (List[Tuple[str | int, str | int]]) – 指定要计算雅可比矩阵的输入-输出对的元组列表。输入和输出可以通过它们在核函数输入/输出参数中出现的整数索引来标识，或通过相应的参数名字符串来标识。如果为 None，则计算所有输入-输出对的雅可比矩阵。

  • device (Device | str | None) – 启动设备（可选）。仅当 function 是 Warp 核函数时使用。

  • max_blocks – 使用的 CUDA 线程块的最大数量。仅当 function 是 Warp 核函数时使用。

  • block_dim – 每个块的线程数。仅当 function 是 Warp 核函数时使用。

  • max_inputs_per_var – 评估输入-输出对雅可比矩阵时考虑的最大输入维度数。如果值 <= 0，则评估所有输入维度。

  • eps (float) – 有限差分步长。

  • plot_jacobians – 如果为 True，则在图中可视化计算出的雅可比矩阵（需要 matplotlib）。

  • metadata (FunctionMetadata) – 核函数的元数据，包含输入和输出标签、步长和数据类型。如果为 None 或为空，则从核函数或函数推断元数据。

  • kernel (Kernel) – 已弃用的参数。请改用 function 参数。

返回：

一个雅可比矩阵字典，其中键是输入和输出索引的元组，值是雅可比矩阵。

返回类型：

Dict[Tuple[int, int], array]

warp.autograd.jacobian_plot#

warp.autograd.jacobian_plot(
jacobians,
kernel,
inputs=None,
outputs=None,
show_plot=True,
show_colorbar=True,
scale_colors_per_submatrix=False,
title=None,
colormap='coolwarm',
log_scale=False,
)[source]#

通过 jacobian() 或 jacobian_fd() 计算的雅可比矩阵可视化为一个组合图像图。需要安装 matplotlib 包。

参数：

  • jacobians (Dict[Tuple[int, int], array]) – 一个雅可比矩阵字典，其中键是输入和输出索引的元组，值是雅可比矩阵。

  • kernel (FunctionMetadata | Kernel) – Warp 核函数，使用 @wp.kernel 装饰器修饰，或具有核函数/函数属性的 FunctionMetadata 实例。

  • inputs (Sequence) – 输入变量列表。

  • outputs (Sequence) – 输出变量列表。已弃用，将在未来的 Warp 版本中移除。

  • show_plot – 如果为 True，则通过 plt.show() 显示图。

  • show_colorbar – 如果为 True，则在图旁边显示颜色条（如果 scale_colors_per_submatrix 为 True，则在每个子矩阵旁边显示颜色条）。

  • scale_colors_per_submatrix – 如果为 True，则分别考虑每个雅可比子矩阵的最小值和最大值进行颜色缩放。否则，使用所有雅可比矩阵的全局最小值和最大值。

  • title (str) – 图的标题（可选）。

  • colormap (str) – 用于图的颜色映射。

  • log_scale – 如果为 True，则对图像图中显示的矩阵值使用对数比例。

返回：

创建的 Matplotlib 图。

示例用法#

import warp as wp
import warp.autograd

@wp.kernel
def my_kernel(
    a: wp.array(dtype=float), b: wp.array(dtype=wp.vec3),
    out1: wp.array(dtype=wp.vec2), out2: wp.array(dtype=wp.quat),
):
    tid = wp.tid()
    ai, bi = a[tid], b[tid]
    out1[tid] = wp.vec2(ai * wp.length(bi), -ai * wp.dot(bi, wp.vec3(0.1, 1.0, -0.1)))
    out2[tid] = wp.normalize(wp.quat(ai, bi[0], bi[1], bi[2]))

a = wp.array([2.0, -1.0], dtype=wp.float32, requires_grad=True)
b = wp.array([wp.vec3(3.0, 1.0, 2.0), wp.vec3(-4.0, -1.0, 0.0)], dtype=wp.vec3, requires_grad=True)
out1 = wp.zeros(2, dtype=wp.vec2, requires_grad=True)
out2 = wp.zeros(2, dtype=wp.quat, requires_grad=True)

# compute the Jacobian matrices for all input/output pairs of the kernel using the autodiff engine
jacs = wp.autograd.jacobian(
    my_kernel, dim=len(a), inputs=[a, b], outputs=[out1, out2],
    plot_jacobians=True)
../_images/kernel_jacobian_ad.svg

jacs 字典包含作为 Warp 数组的雅可比矩阵，对应于核函数的所有输入/输出对。plot_jacobians 参数使用 jacobian_plot() 函数可视化雅可比矩阵。子图显示了每个输入（列）和输出（行）对的雅可比矩阵。这些图像图中的主（粗）网格线分隔了相应 Warp 数组的数组元素。由于核函数参数 b、out1 和 out2 是具有向量类型元素的 Warp 数组，因此相应子图中的次（细、虚线）网格线表示向量元素。

使用 gradcheck() 函数检查梯度精度

passed = wp.autograd.gradcheck(
    my_kernel, dim=len(a), inputs=[a, b], outputs=[out1, out2],
    plot_relative_error=False, plot_absolute_error=False,
    raise_exception=False, show_summary=True)

assert passed

输出

输入

输出

最大绝对误差

最大相对误差

通过

a

out1

1.5134811e-03

4.0449476e-04

 PASS

a

out2

1.1073798e-04

1.4098687e-03

 PASS

b

out1

9.8955631e-04

4.6023726e-03

 PASS

b

out2

3.3494830e-04

1.2789593e-02

 PASS
核函数 my_kernel 的梯度检查通过

除了评估核函数所有输入和输出的雅可比矩阵外，我们还可以将计算限制在特定的输入/输出子集上

jacs = wp.autograd.jacobian(
    my_kernel, dim=len(a), inputs=[a, b], outputs=[out1, out2],
    plot_jacobians=True,
    # select which input/output pairs to compute the Jacobian for
    input_output_mask=[("a", "out1"), ("b", "out2")],
    # limit the number of dimensions to query per output array
    max_outputs_per_var=5,
)
../_images/kernel_jacobian_ad_subset.svg

返回的雅可比矩阵现在被限制在 input_output_mask 参数中指定的输入/输出对。此外，我们使用 max_outputs_per_var 参数将每个输出数组的梯度评估维度数限制为 5。相应的未评估雅可比元素被设置为 NaN。

此外，还可以通过 gradcheck_tape() 函数检查记录在 Tape 上的多个核函数的梯度。在这里，核函数启动的输入和输出用于计算每个核函数启动的雅可比矩阵，并将其与有限差分进行比较

tape = wp.Tape()
with tape:
    wp.launch(my_kernel_1, dim=len(a), inputs=[a, b], outputs=[out1, c])
    wp.launch(my_kernel_2, dim=len(c), inputs=[c], outputs=[out2])

passed = wp.autograd.gradcheck_tape(tape, raise_exception=False, show_summary=True)

assert passed

可视化计算图#

通过自动微分计算梯度可能容易出错，有时数组会缺少 requires_grad 设置，或者在核函数之间传递了错误的数组。为了帮助调试梯度计算，Warp 提供了一个 Tape.visualize() 方法，该方法生成记录在 tape 上的核函数启动的图可视化，格式为 GraphViz dot 格式。可视化显示了 Warp 数组如何用作核函数启动的输入和输出。

示例用法

import warp as wp


@wp.kernel
def add(a: wp.array(dtype=float), b: wp.array(dtype=float), c: wp.array(dtype=float)):
    tid = wp.tid()
    c[tid] = a[tid] + b[tid]


tape = wp.Tape()

a = wp.array([2.0], dtype=wp.float32)
b = wp.array([3.0], dtype=wp.float32, requires_grad=True)
c = wp.array([4.0], dtype=wp.float32)
d = c
e = wp.array([5.0], dtype=wp.float32, requires_grad=True)

result = wp.zeros(1, dtype=wp.float32, requires_grad=True)

with tape:
    wp.launch(add, dim=1, inputs=[b, e], outputs=[a])

    # ScopedTimer registers itself as a scope on the tape
    with wp.ScopedTimer("Adder"):

        # we can also manually record scopes
        tape.record_scope_begin("Custom Scope")
        wp.launch(add, dim=1, inputs=[a, b], outputs=[c])
        tape.record_scope_end()

        wp.launch(add, dim=1, inputs=[d, a], outputs=[result])


tape.visualize(
    filename="tape.dot",
    array_labels={a: "a", b: "b", c: "c", e: "e", result: "result"},
)

这将生成一个 tape.dot 文件，可以使用 GraphViz 工具集进行可视化

dot -Tsvg tape.dot -o tape.svg

生成的 SVG 图像可以在网页浏览器中渲染

../_images/tape.svg

图可视化将核函数启动显示为灰色框，其下方的端口指示输入和输出参数。数组显示为椭圆，其中灰色椭圆表示不需要梯度的数组，绿色椭圆表示具有 requires_grad=True 的数组。

在上面的示例中，我们可以看到数组 c 没有设置 requires_grad 标志，这意味着梯度不会通过此路径传播。

注意

可以使用 tape.visualize() 方法的 array_labels 参数为数组自定义名称。

数组覆盖跟踪#

在计算图中无意中覆盖参与计算图的数组是一个常见错误。例如：

with wp.Tape() as tape:

    # step 1
    wp.launch(compute_forces, dim=n, inputs=[pos0, vel0], outputs=[force])
    wp.launch(simulate, dim=n, inputs=[pos0, vel0, force], outputs=[pos1, vel1])

    # step 2 (error, we are overwriting previous forces)
    wp.launch(compute_forces, dim=n, inputs=[pos1, vel1],  outputs=[force])
    wp.launch(simulate, dim=n, inputs=[pos1, vel1, force], outputs=[pos2, vel2])

    # compute loss
    wp.launch(compute_loss, dim=n, inputs=[pos2], outputs=[loss])

tape.backward(loss)

反向运行 tape 将错误地计算损失相对于 pos0 和 vel0 的梯度，因为 force 在第二个模拟步骤中被覆盖了。force 相对于 pos1 和 vel1 的伴随梯度将是正确的，因为前向传播中存储的 force 值仍然正确，但 force 相对于 pos0 和 vel0 的伴随梯度将是错误的，因为用于此计算的 force 值是在步骤 2 中计算的，而不是步骤 1。解决方案是分配两个 force 数组，force0 和 force1，这样我们就不会覆盖参与计算图的数据。

这类问题归结为一个需要避免的单一模式：在读取数组后写入它。这通常发生在连续的核函数启动之间 (A)，但也可能发生在单个核函数内部 (B)。

A: 核函数间覆盖

@wp.kernel
def square_kernel(x: wp.array(dtype=float), y: wp.array(dtype=float)):
    tid = wp.tid()
    y[tid] = x[tid] * x[tid]


@wp.kernel
def overwrite_kernel(z: wp.array(dtype=float), x: wp.array(dtype=float)):
    tid = wp.tid()
    x[tid] = z[tid]


@wp.kernel
def loss_kernel(x: wp.array(dtype=float), loss: wp.array(dtype=float)):
    tid = wp.tid()
    wp.atomic_add(loss, 0, x[tid])


a = wp.array(np.array([1.0, 2.0, 3.0]), dtype=float, requires_grad=True)
b = wp.zeros_like(a)
c = wp.array(np.array([-1.0, -2.0, -3.0]), dtype=float, requires_grad=True)
loss = wp.zeros(1, dtype=float, requires_grad=True)

tape = wp.Tape()
with tape:
    wp.launch(square_kernel, a.shape, inputs=[a], outputs=[b])
    wp.launch(overwrite_kernel, c.shape, inputs=[c], outputs=[a])
    wp.launch(loss_kernel, a.shape, inputs=[b, loss])

tape.backward(loss)

print(a.grad)

上述代码的输出会产生 a 的不正确梯度

[-2. -4. -6.]

如果移除 overwrite_kernel 启动，我们将获得 a.grad 的正确结果 [2. 4. 6.]。

B: 核函数内覆盖

@wp.kernel
def readwrite_kernel(a: wp.array(dtype=float), b: wp.array(dtype=float)):
    tid = wp.tid()
    b[tid] = a[tid] * a[tid]
    a[tid] = 1.0

@wp.kernel
def loss_kernel(x: wp.array(dtype=float), loss: wp.array(dtype=float)):
    tid = wp.tid()
    wp.atomic_add(loss, 0, x[tid])

a = wp.array(np.array([1.0, 2.0, 3.0]), dtype=float, requires_grad=True)
b = wp.zeros_like(a)
loss = wp.zeros(1, dtype=float, requires_grad=True)

tape = wp.Tape()
with tape:
    wp.launch(readwrite_kernel, dim=a.shape, inputs=[a, b])
    wp.launch(loss_kernel, a.shape, inputs=[b, loss])

tape.backward(loss)

print(a.grad)

上述代码产生不正确的输出

[2. 2. 2.]

如果移除 a[tid] = 1.0，我们将获得 a.grad 的正确结果 [2. 4. 6.]。

如果设置 wp.config.verify_autograd_array_access = True，Warp 将自动检测并报告数组覆盖，包括上述两种情况以及其他有问题配置。它通过在编译期间标记每个核函数中哪些数组参数被读取和/或写入来实现这一点。在运行时，如果一个数组被传递给标记有读取标志的核函数参数，则该数组被标记为已被读取。之后，如果同一个数组被传递给标记有写入标志的核函数参数，则会打印警告（回想一下我们要避免的模式：在*读取*后*写入*）。

注意

设置 wp.config.verify_autograd_array_access = True 将禁用核函数缓存并强制当前模块重建。

注意

此功能尚不支持打包在 Warp struct 中的数组。

如果在图中使用 Tape.record_func()（并提供自己的伴随回调），请务必同时调用 array.mark_write() 和 array.mark_read()，这将手动将您的数组标记为已被写入或读取。

限制和变通方法#

Warp 使用源代码转换方法进行自动微分。在此方法中，反向传播必须记录前向传播期间计算的中间值。这对核函数的可微分性施加了一些限制。

就地数学运算#

原地加法和减法可以在参与反向传播的核函数中使用，例如：

@wp.kernel
def inplace(a: wp.array(dtype=float), b: wp.array(dtype=float)):
    i = wp.tid()

    a[i] -= b[i]


a = wp.full(10, value=10.0, dtype=float, requires_grad=True)
b = wp.full(10, value=2.0, dtype=float, requires_grad=True)

with wp.Tape() as tape:
    wp.launch(inplace, a.shape, inputs=[a, b])

tape.backward(grads={a: wp.ones_like(a)})

print(f"a.grad = {a.grad}")
print(f"b.grad = {b.grad}")

代码产生了预期的输出

a.grad = [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
b.grad = [-1. -1. -1. -1. -1. -1. -1. -1. -1. -1.]

原地乘法和除法不支持，在反向传播中会得到错误的结果。如果wp.config.verbose = True，则在代码生成期间会发出警告。

向量、矩阵和四元数分量赋值#

在核函数内部，对局部定义的向量、矩阵或四元数分量赋值是可微分的，但有一个重要的注意事项：每个分量只能赋值一次（不包括默认初始化）。然后，每个分量可以安全地通过原地加法或减法（+=或-=）操作进行更新，但直接重新赋值（=）将使与该向量、矩阵或四元数相关的梯度计算失效。

动态循环#

目前，动态循环在反向传播中不会被重放或展开，这意味着在循环中计算的、可能对伴随计算必要的中间值不会被更新。

在下面的示例中，计算出了正确的梯度，因为x数组的伴随值不依赖于sum的中间值

@wp.kernel
def dynamic_loop_sum(x: wp.array(dtype=float), loss: wp.array(dtype=float), iters: int):
    sum = float(0.0)

    for i in range(iters):
        sum += x[i]

    wp.atomic_add(loss, 0, sum)


iters = 3
x = wp.full(shape=iters, value=1.0, dtype=float, requires_grad=True)
loss = wp.zeros(1, dtype=float, requires_grad=True)

with wp.Tape() as tape:
    wp.launch(dynamic_loop_sum, dim=1, inputs=[x, loss, iters])

tape.backward(loss)

print(x.grad)

这产生了预期的输出。

[1. 1. 1.]

相反，在此示例中，x数组的伴随值依赖于prod的中间值（adj_x[i] = adj_prod[i+1] * prod[i]），因此梯度计算不正确

@wp.kernel
def dynamic_loop_mult(x: wp.array(dtype=float), loss: wp.array(dtype=float), iters: int):
    prod = float(1.0)

    for i in range(iters):
        prod *= x[i]

    wp.atomic_add(loss, 0, prod)


iters = 3
x = wp.full(shape=iters, value=2.0, dtype=float, requires_grad=True)
loss = wp.zeros(1, dtype=float, requires_grad=True)

with wp.Tape() as tape:
    wp.launch(dynamic_loop_mult, dim=1, inputs=[x, loss, iters])

tape.backward(loss)

print(x.grad)

此代码产生了不正确的梯度，而不是[4. 4. 4.]。

[32.  8.  2.]

我们可以通过切换到静态循环来解决后一种情况（例如，将range(iters)替换为range(3)）。如果循环迭代次数小于或等于在wp.config中设置的max_unroll参数或在模块级别通过wp.set_module_options({"max_unroll": N})设置的参数，静态循环会自动展开，因此循环中的中间值会单独存储。但在不可能这样做的情况下，您可以考虑分配额外的内存来存储动态循环中的中间值。例如，我们可以如下修复上述情况：

@wp.kernel
def dynamic_loop_mult(x: wp.array(dtype=float), prods: wp.array(dtype=float), loss: wp.array(dtype=float), iters: int):
    for i in range(iters):
        prods[i + 1] = x[i] * prods[i]

    wp.atomic_add(loss, 0, prods[iters])


iters = 3
x = wp.full(shape=iters, value=2.0, dtype=float, requires_grad=True)
prods = wp.full(shape=(iters + 1), value=1.0, dtype=float, requires_grad=True)
loss = wp.zeros(1, dtype=float, requires_grad=True)

with wp.Tape() as tape:
    wp.launch(dynamic_loop_mult, dim=1, inputs=[x, prods, loss, iters])

tape.backward(loss)

print(x.grad)

现在，我们得到了预期的梯度。

[4. 4. 4.]

即使数组的伴随值不依赖于动态循环中的中间局部值，局部变量的最终值也可能对伴随计算是必需的。考虑以下情况：

@wp.kernel
def dynamic_loop_sum(x: wp.array(dtype=float), weights: wp.array(dtype=float), loss: wp.array(dtype=float), iters: int):
    sum = float(0.0)
    norm = float(0.0)

    for i in range(iters):
        w = weights[i]
        norm += w
        sum += x[i] * w

    l = sum / norm
    wp.atomic_add(loss, 0, l)


iters = 3
x = wp.full(shape=iters, value=1.0, dtype=float, requires_grad=True)
weights = wp.full(shape=iters, value=1.0, dtype=float, requires_grad=True)
loss = wp.zeros(1, dtype=float, requires_grad=True)

with wp.Tape() as tape:
    wp.launch(dynamic_loop_sum, dim=1, inputs=[x, weights, loss, iters])

tape.backward(loss)

print(x.grad)

此代码产生了不正确的输出。

[inf inf inf]

在反向传播中，当计算用于计算x数组伴随值的sum的伴随值时，发生了除以零的错误：norm在反向传播中没有重新计算，因为动态循环没有重放。这意味着在伴随计算开始时，norm是0.0，而不是在前向传播中计算的值3.0。

对于这种特定情况有另一种补救措施。可以将循环体迁移到一个Warp函数中，从而强制动态循环在反向传播中重放。

@wp.func
def loop(x: wp.array(dtype=float), weights: wp.array(dtype=float), iters: int):
    sum = float(0.0)
    norm = float(0.0)

    for i in range(iters):
        w = weights[i]
        norm += w
        sum += x[i] * w

    return sum, norm


@wp.kernel
def dynamic_loop_sum(x: wp.array(dtype=float), weights: wp.array(dtype=float), loss: wp.array(dtype=float), iters: int):
    sum, norm = loop(x, weights, iters)

    l = sum / norm
    wp.atomic_add(loss, 0, l)


iters = 3
x = wp.full(shape=iters, value=1.0, dtype=float, requires_grad=True)
weights = wp.full(shape=iters, value=0.5, dtype=float, requires_grad=True)
loss = wp.zeros(1, dtype=float, requires_grad=True)

with wp.Tape() as tape:
    wp.launch(dynamic_loop_sum, dim=1, inputs=[x, weights, loss, iters])

tape.backward(loss)

print(x.grad)

现在，上面的代码产生了预期的结果。

[0.33333334 0.33333334 0.33333334]

然而，这之所以奏效，是因为x数组的伴随值不需要sum的中间值；它们只需要sum的伴随值。通常，此变通方法仅适用于简单的加法/减法操作，例如+=或-=。

注意

在后续版本中，我们将允许用户在某些情况下强制展开动态循环，从而避免这些变通方法。