cuquantum.Network

class cuquantum.Network(subscripts, *operands, qualifiers=None, options=None, stream=None)[源代码]

创建一个张量网络对象，该对象指定为爱因斯坦求和表达式。

爱因斯坦求和约定提供了一种表示许多张量网络操作的优雅方式。此对象允许用户投入大量精力来计算最佳收缩路径，以及预先自动调整收缩，以便对同一网络拓扑（具有相同爱因斯坦求和表达式的不同输入张量或“操作数”）进行重复收缩。另请参阅 contract_path() 和 autotune()。

对于爱因斯坦求和表达式，显式形式和隐式形式都受支持。

在隐式形式中，输出模式标签从求和表达式中推断出来，并按字典顺序重新排序。例如表达式 'ij,jh'，其输出模式标签为 'hi'。（这对应于矩阵乘法，后跟转置。）

在显式形式中，输出模式标签可以直接在求和表达式中使用标识符 '->' 后声明。例如表达式 'ij,jh->ih'（对应于矩阵乘法）。

要指定爱因斯坦求和表达式，支持下标格式（如上所示）和交错格式。

交错格式是指定操作数及其模式标签的另一种方式，格式为 Network(op0, modes0, op1, modes1, ..., [modes_out])，其中 opN 是第 N 个操作数，modesN 是可哈希和可比较对象（字符串、整数等）的序列，表示第 N 个操作数的模式标签。

支持省略号广播。

可以通过将 logging.Logger 对象传递给 NetworkOptions 或通过在默认使用的根 logger 对象中设置适当的选项来获得有关网络上各种操作的其他信息

>>> import logging
>>> logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)-8s %(message)s', datefmt='%m-%d %H:%M:%S')
参数

  • subscripts – 定义爱因斯坦求和表达式的模式标签（下标），以逗号分隔的字符序列形式表示。表达式中允许使用 Unicode 字符，从而扩展了可以使用爱因斯坦求和约定指定的张量网络的大小。

  • operands – 张量（类 ndarray 对象）序列。当前支持的类型有 numpy.ndarray、cupy.ndarray 和 torch.Tensor。

  • qualifiers – 将张量限定符指定为 numpy.ndarray，其中包含长度等于操作数数量的 tensor_qualifiers_dtype 对象。

  • options – 将张量网络的选项指定为 NetworkOptions 对象。或者，也可以提供包含 NetworkOptions 构造函数参数的 dict。如果未指定，则该值将设置为默认构造的 NetworkOptions 对象。

  • stream – 提供用于网络构建的 CUDA 流，这对于流序操作（如分配内存）是必需的。可接受的输入包括 cudaStream_t（作为 Python int）、cupy.cuda.Stream 和 torch.cuda.Stream。如果未提供流，则将使用当前流。

另请参阅

contract_path()、autotune()、contract()、reset_operands()

示例

>>> from cuquantum import Network
>>> import numpy as np

定义张量网络的参数

>>> expr = 'ehl,gj,edhg,bif,d,c,k,iklj,cf,a->ba'
>>> shapes = [(8, 2, 5), (5, 7), (8, 8, 2, 5), (8, 6, 3), (8,), (6,), (5,), (6, 5, 5, 7), (6, 3), (3,)]

使用 NumPy 创建输入张量

>>> operands = [np.random.rand(*shape) for shape in shapes]

创建 Network 对象

>>> tn = Network(expr, *operands)

查找最佳收缩顺序

>>> path, info = tn.contract_path({'samples': 500})

自动调优网络

>>> tn.autotune(iterations=5)

执行收缩。结果类型和设备与操作数相同

>>> r1 = tn.contract()

将操作数重置为新值

>>> operands = [i*operand for i, operand in enumerate(operands, start=1)]
>>> tn.reset_operands(*operands)

获取新收缩的结果

>>> r2 = tn.contract()
>>> from math import factorial
>>> np.allclose(r2, factorial(len(operands))*r1)
True

最后，释放网络资源。如果不进行此调用，则可能会阻碍进一步的操作（特别是当网络很大时），因为内存只有在对象超出作用域时才会被释放。（为了避免必须显式进行此调用，建议将 Network 对象用作上下文管理器。）

>>> tn.free()

如果操作数在 GPU 上，也可以使用就地操作来更新它们。在这种情况下，可以跳过对 reset_operands() 的调用 – 后续的 contract() 调用将使用相同的操作数（具有更新的内容）。以下示例使用 CuPy 操作数来说明这一点，并演示了 Network 上下文的用法（以便跳过调用 free()）

>>> import cupy as cp
>>> expr = 'ehl,gj,edhg,bif,d,c,k,iklj,cf,a->ba'
>>> shapes = [(8, 2, 5), (5, 7), (8, 8, 2, 5), (8, 6, 3), (8,), (6,), (5,), (6, 5, 5, 7), (6, 3), (3,)]
>>> operands = [cp.random.rand(*shape) for shape in shapes]
>>>
>>> with Network(expr, *operands) as tn:
...     path, info = tn.contract_path({'samples': 500})
...     tn.autotune(iterations=5)
...
...     # Perform the contraction
...     r1 = tn.contract()
...
...     # Update the operands in place
...     for i, operand in enumerate(operands, start=1):
...         operand *= i
...
...     # Perform the contraction with the updated operand values
...     r2 = tn.contract()
...
... # The resources used by the network are automatically released when the context ends.
>>>
>>> from math import factorial
>>> cp.allclose(r2, factorial(len(operands))*r1)
array(True)

PyTorch CPU 和 GPU 张量可以以相同的方式作为输入操作数传递。

要计算网络关于输入操作数 (NumPy/CuPy/PyTorch) 的梯度，可以使用 gradients() 方法。要启用梯度计算，应执行以下操作

  1. 使用 qualifiers 参数创建网络

  2. 在 gradients() 方法之前调用 contract() 方法

  3. 使用输出梯度为 gradients() 方法播种（请参阅文档了解要求）

以下是一个最小示例

>>> from cuquantum import cutensornet as cutn
>>> expr = "ijk,jkl,klm,lmn"
>>> shapes = ((3, 4, 5), (4, 5, 3), (5, 3, 2), (3, 2, 6))
>>> operands = [cp.random.rand(*shape) for shape in shapes]
>>> qualifiers = np.zeros(len(shapes), dtype=cutn.tensor_qualifiers_dtype)
>>> qualifiers[:]["requires_gradient"] = 1  # request gradients for all input tensors
>>>
>>> with Network(expr, *operands, qualifiers=qualifiers) as tn:
...     path, info = tn.contract_path()
...
...     # Perform the contraction
...     r = tn.contract()
...
...     # Perform the backprop
...     input_grads = tn.gradients(cp.ones_like(r))
...
>>>

对于设置了 requires_grad 属性的 PyTorch CPU/GPU 张量，无需传递 qualifiers 参数。请注意，此 Network 类及其方法不是 PyTorch 算符，并且不会向 PyTorch 的 autograd 图添加任何节点。对于原生的、可微分的 PyTorch 算符，请使用 cuquantum.contract() 函数。

有关指定爱因斯坦求和表达式以及指定张量网络和优化器选项的更多示例，请参阅 contract()。

方法

__init__(subscripts, *operands, qualifiers=None, options=None, stream=None)[源代码]
autotune(*, iterations=3, stream=None, release_workspace=False)[源代码]

自动调优网络以降低收缩成本。

如果 Network 对象用于执行多次收缩，则建议执行此可选步骤。

参数

  • iterations – 自动调优的迭代次数。请参阅 CUTENSORNET_CONTRACTION_AUTOTUNE_MAX_ITERATIONS。

  • stream – 提供 CUDA 流以用于自动调优操作。可接受的输入包括 cudaStream_t（作为 Python int）、cupy.cuda.Stream 和 torch.cuda.Stream。如果未提供流，则将使用当前流。

  • release_workspace – True 值指定 Network 对象应在函数返回时将工作区内存释放回包内存池，而 False 值指定 Network 对象应保留内存。如果应用程序在连续调用（相同或不同的）执行 API（例如 autotune()、contract() 或 gradients()）之间执行其他大量消耗内存的操作，则可以将此选项设置为 True。但这会在每次调用时因从包内存池获取和释放工作区内存而产生少量开销。默认值为 False。

contract(*, slices=None, stream=None, release_workspace=False)[source]

收缩网络并返回结果。

参数

  • slices – 指定要收缩的切片，可以使用 Python range 表示连续的切片 ID，或者使用 Python 序列对象表示任意切片 ID。如果未指定，则将收缩所有切片。

  • stream – 提供 CUDA 流以用于收缩操作。可接受的输入包括 cudaStream_t（作为 Python int）、cupy.cuda.Stream 和 torch.cuda.Stream。如果未提供流，则将使用当前流。

  • release_workspace – True 值指定 Network 对象应在函数返回时将工作区内存释放回包内存池，而 False 值指定 Network 对象应保留内存。如果应用程序在连续调用（相同或不同的）执行 API（例如 autotune()、contract() 或 gradients()）之间执行其他大量消耗内存的操作，则可以将此选项设置为 True。但这会在每次调用时因从包内存池获取和释放工作区内存而产生少量开销。默认值为 False。

返回值

结果类型和设备与操作数相同。

contract_path(optimize=None)[source]

计算最佳收缩路径以及确保在指定内存限制内执行收缩所需的任何切片。

参数

optimize – 此参数将路径优化选项指定为 OptimizerOptions 对象。或者，也可以提供包含 OptimizerOptions 构造函数参数的字典。如果未指定，则该值将设置为默认构造的 OptimizerOptions 对象。

返回值

一个 2 元组 (path, opt_info)

  • path : 一系列操作数序号对，表示 numpy.einsum_path() 格式的最佳收缩顺序。

  • opt_info : OptimizerInfo 类型的对象，包含有关最佳收缩顺序的信息。

返回类型

元组

注释

  • 如果提供了路径，则如果需要切片，用户还必须设置切片模式。

free()[source]

释放网络资源。

建议在上下文中使用 Network 对象，但如果不可能，则必须显式调用此方法以确保正确清理网络资源。

gradients(output_gradient, *, stream=None, release_workspace=False)[source]

计算网络梯度（相对于需要梯度的输入操作数）。

在调用此方法之前，必须已执行完整收缩（通过调用 contract()），否则会引发错误。

参数

  • output_gradient – 张量，其包（NumPy/CuPy/PyTorch）、形状、dtype、步幅和位置（CPU/GPU）与收缩输出（由 contract() 返回）相同，而收缩输出又与输入操作数共享相同的属性。在链式法则设置中，output_gradient 是相对于输出张量的梯度。

  • stream – 提供 CUDA 流以用于梯度计算。可接受的输入包括 cudaStream_t（作为 Python int）、cupy.cuda.Stream 和 torch.cuda.Stream。如果未提供流，则将使用当前流。

  • release_workspace – True 值指定 Network 对象应在函数返回时将工作区内存释放回包内存池，而 False 值指定 Network 对象应保留内存。如果应用程序在连续调用（相同或不同的）执行 API（例如 autotune()、contract() 或 gradients()）之间执行其他大量消耗内存的操作，则可以将此选项设置为 True。但这会在每次调用时因从包内存池获取和释放工作区内存而产生少量开销。默认值为 False。

返回值

梯度张量序列。结果的长度和类型与输入操作数相同，并且在同一设备上。对于未请求的梯度分量，将返回 None。

注意

对于 PyTorch 操作数，调用此方法不会被 autograd 图跟踪。

警告

此 API 是实验性的，将来可能会发生更改。

reset_operands(*operands, stream=None)[source]

重置此 Network 实例持有的操作数。

此方法有两种用例：（1）当

原始操作数位于 CPU 上时，它可以用于为执行提供新的操作数；或者（2）可以通过为 operands 参数传递 None 来释放对先前操作数的内部引用，并使其内存可供其他用途使用。在这种情况下，必须再次调用此方法，以在再次调用执行 API（如 autotune()、contract() 或 gradients()）之前提供所需的操作数。

当操作数驻留在 GPU 上并且使用就地操作更新操作数值时，不需要此方法。

此方法将对新操作数执行各种检查，以确保

  • 形状、步幅、数据类型与旧操作数匹配。

  • 操作数所属的包与旧操作数匹配。

  • 如果输入张量在 GPU 上，则库包和设备必须匹配。

参数

  • operands – 请参阅 Network 的文档。

  • stream – 提供 CUDA 流以用于重置操作数（如果操作数在 CPU 上提供，则用于将操作数复制到 GPU）。可接受的输入包括 cudaStream_t（作为 Python int）、cupy.cuda.Stream 和 torch.cuda.Stream。如果未提供流，则将使用当前流。