PyTorch中如何使用matmul完成2D置换矩阵与3D张量乘法运算

直接调用torch.matmul(P, T)不符合PyTorch的高维矩阵乘法维度匹配规则：
matmul执行运算时，默认取参与张量的最后两个维度作为矩阵运算维度，前面的维度视为批量维度。你定义的T维度顺序为(C, W, H)=(2,2,2)，直接和形状为(2,2)的P做乘法时，matmul会自动匹配P的最后一维2和T的倒数第二维2运算，实际是对W维度做变换，而非你需要的C维度。

以下两种方法都可以用matmul完成你需要的通道置换运算： #

方法1：展平空间维度后运算 #

把T的后两个空间维度（W、H）展平为一维，乘完后再恢复原始形状即可：

import torch

P = torch.tensor([[0,1],[1,0]], dtype=torch.float32)
M1 = torch.tensor([[1,2],[3,4]], dtype=torch.float32)
M2 = torch.tensor([[5,6],[7,8]], dtype=torch.float32)
T = torch.stack([M1, M2], dim=0) # 形状为 (2,2,2) 即 (C,W,H)

# 展平后两维得到形状 (C, W*H) = (2,4)
T_flat = T.flatten(start_dim=1)
# 矩阵乘法后reshape回原始维度
PT = torch.matmul(P, T_flat).reshape(T.shape)

方法2：调整维度顺序后运算 #

把需要参与矩阵运算的C维度挪到T的最后一位，运算完成后再换回原维度顺序：

# 调整维度顺序为 (W, H, C) = (2,2,2)
T_permuted = T.permute(1,2,0)
# 乘法完成后调整回 (C,W,H) 的原始维度顺序
PT = torch.matmul(P, T_permuted).permute(2,0,1)

两种方法最终输出的PT均为你需要的[[[5,6],[7,8]], [[1,2],[3,4]]]结果。

内容的提问来源于stack exchange，提问作者phil