PyTorch中用torch.block_diag()将多个矩阵沿对角线拼接成一个大矩阵的函数

torch.block_diag()是PyTorch中用于将多个矩阵沿对角线拼接成一个大矩阵的函数。这个函数可以用于构建卷积神经网络中的卷积核矩阵，或者构建变分自编码器等需要对多个线性变换进行堆叠的模型。

torch.block_diag()函数的语法如下：

torch.block_diag(*args)
1

其中，*args是要拼接的矩阵，可以是一个或多个Tensor对象。

下面是一个简单的代码示例，演示了如何使用torch.block_diag()函数将三个矩阵沿对角线拼接成一个大矩阵：

import torch

a = torch.tensor([[1, 2], [3, 4]])
b = torch.tensor([[5, 6]])
c = torch.tensor([[7], [8]])

result = torch.block_diag(a, b, c)
print(result)
1
2
3
4
5
6
7
8

输出结果为：

tensor([[1, 2, 0, 0],
        [3, 4, 0, 0],
        [0, 0, 5, 6],
        [0, 0, 0, 7],
        [0, 0, 0, 8]])
1
2
3
4
5

在上面的例子中，我们定义了三个矩阵a、b、c，分别为2×2、1×2、2×1的矩阵。使用torch.block_diag()函数将这三个矩阵沿对角线拼接成一个5×4的矩阵，其中未被填充的部分用0填充。

注意，在使用torch.block_diag()函数时，传入的矩阵应该具有相同的数据类型和设备类型。如果有一个矩阵的类型不一致，那么会抛出类型不匹配的异常。如果要将一个CPU上的矩阵和一个GPU上的矩阵拼接在一起，需要先使用.to()方法将它们转换成同一设备类型。

import torch

a = torch.tensor([[1, 2], [3, 4]], device='cuda')
b = torch.tensor([[5, 6]], device='cuda')
c = torch.tensor([[7], [8]])

# TypeError: block_diag(): Expected all inputs to be on the same device, but found at least two devices, cuda:0 and cpu!
result = torch.block_diag(a, b, c)
1
2
3
4
5
6
7
8

上面的代码会抛出TypeError异常，因为c矩阵在CPU上，而a和b矩阵在GPU上。需要将c矩阵转换为GPU上的Tensor，才能与a和b矩阵拼接在一起：

import torch

a = torch.tensor([[1, 2], [3, 4]], device='cuda')
b = torch.tensor([[5, 6]], device='cuda')
c = torch.tensor([[7], [8]], device='cuda')
1
2
3
4
5