python - 朴素的平铺矩阵乘法与 Python

Question

我正在尝试编写一个函数来对两个大矩阵执行矩阵乘法运算，这两个大矩阵被分成较小的矩阵（瓦片），我很难概念化一个有效的方法。

在下面的脚本中，两个张量/矩阵a和b被分割成多个形状为 的子张量/矩阵（瓦片）tile_shape。鉴于这些图块，我想执行与 a.dot(b) 给出相同结果的操作，即向量矩阵乘法。

为简单起见，瓦片都是相同的形状，并且在给定张量的长度/宽度不能被 整除的情况下填充为零tile shape。我希望该方法是模块化的，以便可以指定任何给定的有效输入形状和平铺形状。任何帮助是极大的赞赏！

import numpy as np
import math

class Tiles:
    def __init__(self, tiles, tensor_shape):
        self.tiles = tiles
        self.tile_shape = tiles.shape[-2:]
        self.tensor_shape = tensor_shape

def gen_tiles(tensor, tile_shape):
    assert len(tensor.shape) == 2
    tiles = np.empty((math.ceil(tensor.shape[0] / tile_shape[0]), math.ceil(tensor.shape[1] / tile_shape[1]), tile_shape[0], tile_shape[1]))
    for i in range(0, tiles.shape[0]):
        for j in range(0, tiles.shape[1]):
            tile = tensor[i * tile_shape[0]:max((i + 1) * tile_shape[0], tile_shape[0] - 1), j * tile_shape[1]:max((j + 1) * tile_shape[1], tile_shape[1] - 1)]
            padded_tile = np.zeros(shape=tile_shape)
            padded_tile[:tile.shape[0],:tile.shape[1]] = tile
            tiles[i][j][:][:] = padded_tile

    return Tiles(tiles, tensor.shape)

def matmul_tiles(a, b):
    assert a.tensor_shape[1] == b.tensor_shape[0] and a.tile_shape == b.tile_shape
    result = np.zeros((a.tensor_shape[0], b.tensor_shape[1]))
    print(a.tiles.shape[0:2])
    print(b.tiles.shape[0:2])
    # Compute matrix multiplication using tiles

tile_shape = (2, 2)
a = np.random.randint(10, size=(2, 2))
b = np.random.randint(10, size=(2, 3))
a_tiles = gen_tiles(a, tile_shape)
b_tiles = gen_tiles(b, tile_shape)
product = a.dot(b)
tile_product = matmul_tiles(a_tiles, b_tiles)
assert np.isclose(product, tile_product)

标签： pythonnumpymatrix-multiplication