tensorflow 或 pytorch 中的分区矩阵乘法

Question

假设我有大小为 [4, 4] 的矩阵 P，它将（块）划分为 4 个较小的矩阵 [2,2]。我怎样才能有效地将这个块矩阵乘以另一个矩阵（不是分区矩阵但更小）？

让我们假设我们的原始矩阵是：

P = [ 1 1 2 2
      1 1 2 2
      3 3 4 4
      3 3 4 4]

其中拆分为子矩阵：

P_1 = [1 1    , P_2 = [2 2  , P_3 = [3 3   P_4 = [4 4
       1 1]            2 2]          3 3]         4 4]

现在我们的P是：

P = [P_1 P_2
     P_3 p_4]

在下一步中，我想在 P 和较小的矩阵之间进行逐元素乘法，其大小等于子矩阵的数量：

P * [ 1 0   =   [P_1  0  = [1 1 0 0 
      0 0 ]      0    0]    1 1 0 0
                            0 0 0 0
                            0 0 0 0]

Answer 1

不知道有效的方法，但你可以试试这些：

方法一：

使用torch.cat()

import torch

def multiply(a, b):
    x1 = a[0:2, 0:2]*b[0,0]
    x2 = a[0:2, 2:]*b[0,1]
    x3 = a[2:, 0:2]*b[1,0]
    x4 = a[2:, 2:]*b[1,1]
    return torch.cat((torch.cat((x1, x2), 1), torch.cat((x3, x4), 1)), 0)

a = torch.tensor([[1, 1, 2, 2],[1, 1, 2, 2],[3, 3, 4, 4,],[3, 3, 4, 4]])
b = torch.tensor([[1, 0],[0, 0]])
print(multiply(a, b))

输出：

tensor([[1, 1, 0, 0],
        [1, 1, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]])

方法二：

使用torch.nn.functional.pad()

import torch.nn.functional as F
import torch

def multiply(a, b):
    b = F.pad(input=b, pad=(1, 1, 1, 1), mode='constant', value=0)
    b[0,0] = 1
    b[0,1] = 1
    b[1,0] = 1
    return a*b

a = torch.tensor([[1, 1, 2, 2],[1, 1, 2, 2],[3, 3, 4, 4,],[3, 3, 4, 4]])
b = torch.tensor([[1, 0],[0, 0]])
print(multiply(a, b))

输出：

tensor([[1, 1, 0, 0],
        [1, 1, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]])

Answer 2

虽然其他答案可能是解决方案，但这不是一种有效的方法。我想出了另一种方法来解决这个问题（但仍然不完美）。当我们的输入是 3 或 4 维时，以下实现需要太多内存。例如，对于 20*75*1024*1024 的输入大小，以下计算需要大约 12gb ram。

这是我的实现：

import tensorflow as tf

tf.enable_eager_execution()


inps = tf.constant([
    [1, 1, 1, 1, 2, 2, 2, 2],
    [1, 1, 1, 1, 2, 2, 2, 2],
    [1, 1, 1, 1, 2, 2, 2, 2],
    [1, 1, 1, 1, 2, 2, 2, 2],
    [3, 3, 3, 3, 4, 4, 4, 4],
    [3, 3, 3, 3, 4, 4, 4, 4],
    [3, 3, 3, 3, 4, 4, 4, 4],
    [3, 3, 3, 3, 4, 4, 4, 4]])

on_cells = tf.constant([[1, 0, 0, 1]])

on_cells = tf.expand_dims(on_cells, axis=-1)

# replicate the value to block-size (4*4)
on_cells = tf.tile(on_cells, [1, 1, 4 * 4])

# reshape to a format for permutation
on_cells = tf.reshape(on_cells, (1, 2, 2, 4, 4))

# permutation
on_cells = tf.transpose(on_cells, [0, 1, 3, 2, 4])

# reshape
on_cells = tf.reshape(on_cells, [1, 8, 8])

# element-wise operation
print(inps * on_cells)

输出：

tf.Tensor(
[[[1 1 1 1 0 0 0 0]
  [1 1 1 1 0 0 0 0]
  [1 1 1 1 0 0 0 0]
  [1 1 1 1 0 0 0 0]
  [0 0 0 0 4 4 4 4]
  [0 0 0 0 4 4 4 4]
  [0 0 0 0 4 4 4 4]
  [0 0 0 0 4 4 4 4]]], shape=(1, 8, 8), dtype=int32)

Answer 3

您可以考虑以更有效的方式表示您的大块矩阵。

例如，块矩阵

P = [ 1 1 2 2
      1 1 2 2
      3 3 4 4
      3 3 4 4]

可以用

表示

a = [ 1 0    b = [ 1 1 0 0    p = [ 1 2
      1 0          0 0 1 1 ]        3 4 ]
      0 1
      0 1 ]

作为

P = a @ p @ b

与（@表示矩阵相乘）。矩阵a和brepresents/encode分块结构P和小p表示每个分块的值

现在，如果您想将 (element-wise) p 与一个小的 (2x2) 矩阵 q 相乘，您只需

a @ (p * q) @ b

一个简单的pytorch例子

In [1]: a = torch.tensor([[1., 0], [1., 0], [0., 1], [0, 1]])
In [2]: b = torch.tensor([[1., 1., 0, 0], [0, 0, 1., 1]]) 
In [3]: p=torch.tensor([[1., 2.], [3., 4.]])
In [4]: q = torch.tensor([[1., 0], [0., 0]])

In [5]: a @ p @ b

Out[5]:
tensor([[1., 1., 2., 2.],
        [1., 1., 2., 2.],
        [3., 3., 4., 4.],
        [3., 3., 4., 4.]])

In [6]: a @ (p*q) @ b

Out[6]:
tensor([[1., 1., 0., 0.],
        [1., 1., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])

我把它留给你作为练习如何在给定块大小的情况下有效地生成 "structure" 矩阵 a 和 b。

Answer 4

以下是通用 Tensorflow-based 解决方案，适用于任意形状的输入矩阵 p（大）和 m（小）只要 [=14= 的大小] 在两个轴上都可以被 m 的大小整除。

def block_mul(p, m):
   p_x, p_y = p.shape
   m_x, m_y = m.shape
   m_4d = tf.reshape(m, (m_x, 1, m_y, 1))
   m_broadcasted = tf.broadcast_to(m_4d, (m_x, p_x // m_x, m_y, p_y // m_y))
   mp = tf.reshape(m_broadcasted, (p_x, p_y))
   return p * mp

测试：

import tensorflow as tf

tf.enable_eager_execution()

p = tf.reshape(tf.constant(range(36)), (6, 6))
m = tf.reshape(tf.constant(range(9)), (3, 3))
print(f"p:\n{p}\n")
print(f"m:\n{m}\n")
print(f"block_mul(p, m):\n{block_mul(p, m)}")

输出（Python 3.7.3，Tensorflow 1.13.1）：

p:
[[ 0  1  2  3  4  5]
 [ 6  7  8  9 10 11]
 [12 13 14 15 16 17]
 [18 19 20 21 22 23]
 [24 25 26 27 28 29]
 [30 31 32 33 34 35]]

m:
[[0 1 2]
 [3 4 5]
 [6 7 8]]

block_mul(p, m):
[[  0   0   2   3   8  10]
 [  0   0   8   9  20  22]
 [ 36  39  56  60  80  85]
 [ 54  57  80  84 110 115]
 [144 150 182 189 224 232]
 [180 186 224 231 272 280]]

另一种使用隐式广播的解决方案如下：

def block_mul2(p, m):
   p_x, p_y = p.shape
   m_x, m_y = m.shape
   p_4d = tf.reshape(p, (m_x, p_x // m_x, m_y, p_y // m_y))
   m_4d = tf.reshape(m, (m_x, 1, m_y, 1))
   return tf.reshape(p_4d * m_4d, (p_x, p_y))

Answer 5

如果矩阵很小，您可能可以使用 cat 或 pad。分解的解决方案非常优雅，就像 block_mul 实现的解决方案一样。

另一种解决方案是在 3D 体积中转动 2D 块矩阵，其中每个 2D 切片都是一个块 (P_1, P_2, P_3, P_4)。然后使用广播的力量将每个 2D 切片乘以一个标量。最后重塑输出。重塑不是立竿见影的，但它是可行的，从 numpy 移植到

的 pytorch

在 Pytorch 中：

import torch

h = w = 4
x = torch.ones(h, w)
x[:2, 2:] = 2
x[2:, :2] = 3
x[2:, 2:] = 4

# number of blocks along x and y
nrows=2
ncols=2

vol3d = x.reshape(h//nrows, nrows, -1, ncols)
vol3d = vol3d.permute(0, 2, 1, 3).reshape(-1, nrows, ncols)

out = vol3d * torch.Tensor([1, 0, 0, 0])[:, None, None].float()

# reshape to original
n, nrows, ncols = out.shape
out = out.reshape(h//nrows, -1, nrows, ncols)
out = out.permute(0, 2, 1, 3)
out = out.reshape(h, w)

print(out)

tensor([[1., 1., 0., 0.],
        [1., 1., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])

我没有将其与其他基准进行对比，但这不会像填充那样消耗额外的内存，也不会像连接那样执行缓慢的操作。它还具有易于理解和可视化的优点。

您可以通过 h, w, nrows, ncols 将其推广到任何情况。

tensorflow 或 pytorch 中的分区矩阵乘法

Partitioned matrix multiplication in tensorflow or pytorch

matrix

multiplication

matrix-multiplication

tensorflow

pytorch