TensorNetwork笔记(1)

系列文章目录

---

安装

pip3 install tensornetwork

一、Node 节点

节点是张量网络的基本构建块之一。它们代表计算中的张量。每个轴都有一条对应的边，可以将不同的节点（甚至同一个节点）连接在一起。边数表示基础张量的阶数。例如，没有任何边的节点是标量，有一条边的节点是向量，等等

import tensornetwork as tn
import numpy as np

node = tn.Node(np.eye(2))

二、Edge 边

边描述了张量网络中底层张量的不同收缩。边缘与收缩中涉及的张量轴相关联。张量网络中有 3 种基本类型的边：

标准边 Standard Edges

标准边与您在无向图中找到的任何其他边一样。它们连接 2 个不同的节点并表示相关向量空间之间的点积。在 numpy 术语中，这条边定义了给定轴上的张量点操作。

跟踪边 Trace Edges

跟踪边将节点连接到自身。收缩这种类型的边意味着获取由两个相关轴形成的矩阵的轨。

悬垂边 Dangling Edge

悬垂边是只有一侧指向节点的边，另一侧“悬空”。这些边可以表示输出轴，也可以表示尚未连接到其他悬垂边的中间轴。这些边是在将节点添加到网络时自动创建的。

a = tn.Node(np.eye(2))
b = tn.Node(np.eye(2))
c = tn.Node(np.eye(2))
#悬垂边在节点创建时自动创建。
#我们可以通过这种方式访问它们。
dangling_edge = a.get_edge(1)
dangling_edge = a[1]
#通过将任意两个单独的节点连接在一起来创建标准边。
#我们通过“连接”两条悬挂的边来创建一条新边。
standard_edge = a[0] ^ b[0] # 也以可用 tn.connect(a[0], b[0])
#通过将节点连接到自身来创建跟踪边。
trace_edge = c[0] ^ c[1]

三、张量收缩

tn.contractors有几种可用的收缩算法。

optimal：通过蛮力找到真正的最佳路径。对于超过 10 个节点，它可能会非常慢。
greedy: 不断地做尽可能便宜的收缩。适用于具有许多节点的网络的默认设置。
branch：蛮力搜索，但只检查n每一步的最高可能性。
auto：自动决定使用上述 3 种算法中的哪一种。
以greedy为例

# Here, we will contract the following shaped network.
# O - O
# | X |
# O - O
a = tn.Node(np.ones((2, 2, 2)))
b = tn.Node(np.ones((2, 2, 2)))
c = tn.Node(np.ones((2, 2, 2)))
d = tn.Node(np.ones((2, 2, 2)))
# Make the network fully connected.
a[0] ^ b[0]
a[1] ^ c[1]
a[2] ^ d[2]
b[1] ^ d[1]
b[2] ^ c[2]
c[0] ^ d[0]
# We are using the "greedy" contraction algorithm.
# Other algorithms we support include "optimal" and "branch".

# Finding the optimial contraction order in the general case is NP-Hard,
# so there is no single algorithm that will work for every tensor network.
# However, there are certain kinds of networks that have nice properties that
# we can expliot to making finding a good contraction order easier.
# These types of contraction algorithms are in developement, and we welcome 
# PRs!

# `tn.reachable` will do a BFS to get all of the nodes reachable from a given
# node or set of nodes.
# nodes = {a, b, c, d}
nodes = tn.reachable(a)
result = tn.contractors.greedy(nodes)
print(result.tensor)

输出为64

四、优化收缩

在计算的中间状态，两个节点有多个连接它们的边是很常见的。如果这些边中只有一个被收缩，那么所有剩余的边都成为跟踪边。这通常是非常低效的，因为新节点将分配比最终所需更多的内存。由于迹线边缘仅对基础矩阵的对角线求和，因此在第一次收缩期间计算的所有其他值都是无用的。同时收缩所有这些边总是更有效。
方法contract_between或contract_parallel将自动为您执行此操作。将这些方法与一次收缩一个优势进行比较时，您应该会看到巨大的加速。

def one_edge_at_a_time(a, b):
  node1 = tn.Node(a)
  node2 = tn.Node(b)
  edge1 = node1[0] ^ node2[0]
  edge2 = node1[1] ^ node2[1]
  tn.contract(edge1)
  result = tn.contract(edge2)
  return result.tensor

def use_contract_between(a, b):
  node1 = tn.Node(a)
  node2 = tn.Node(b)
  node1[0] ^ node2[0]
  node1[1] ^ node2[1]
  # This is the same as 
  # tn.contract_between(node1, node2)
  result = node1 @ node2
  return result.tensor

a = np.ones((1000, 1000))
b = np.ones((1000, 1000))
print(one_edge_at_a_time(a,b))
print("Running one_edge_at_a_time")
%timeit one_edge_at_a_time(a, b)
print("Running use_cotract_between")
%timeit use_contract_between(a, b)

1000000.0
Running one_edge_at_a_time
7.73 ms ± 441 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Running use_cotract_between
394 µs ± 20 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)