L有一個形狀的單一來源（正方形）矩陣(N, N)
import torch as pt
import numpy as np
N = 4
L = pt.arange(N*N).reshape(N, N) # or np.arange(N*N).reshape(N, N)
L = tensor([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11],
[12, 13, 14, 15]])
m和一個形狀的布爾掩碼矩陣（向量的向量）(K, N)，我想根據它從中提取子矩陣L。
K = 3
m = tensor([[ True, True, False, False],
[False, True, True, False],
[False, True, False, True]])
L[m[i]][:, m[i]]我知道如何通過呼叫any來使用單個掩碼向量提取單個子矩陣i。因此，例如，對于i=0，我們會得到
tensor([[ 0, 1],
[ 4, 5]])
但我需要沿整個“批次”維度執行操作。我正在尋找的最終結果可以通過
res = []
for i in range(K):
res.append(L[m[i]][:, m[i]])
output = pt.stack(res)
但是，我希望有一個更好的解決方案，不包括 for 回圈。我意識到如果m沿最后一個維度 ( dim/axis=1) 的總和不是恒定的，for 回圈解決方案本身就會崩潰，但如果我可以保證它是恒定的，是否有更好的解決方案？如果沒有，更改選擇器表示會有幫助嗎？為了方便起見，我選擇了布爾掩碼，但我更喜歡更好的性能。
uj5u.com熱心網友回復：
請注意，您可以通過索引和廣播一起獲得第一個方塊：
r = torch.tensor([0,1])
L[r[:,None], r]
輸出：
tensor([[0, 1],
[4, 5]])
同樣的原理可以應用于第二個正方形：
r = torch.tensor([1,2])
L[r[:,None], r]
輸出：
tensor([[ 5, 6],
[ 9, 10]])
結合起來你得到：
i = torch.tensor([[0, 1], [1, 2]])
L[i[:,:,None], i[:,None]]
輸出：
tensor([[[ 0, 4],
[ 1, 5]],
[[ 5, 9],
[ 6, 10]]])
所有 3 個方格：
i = torch.tensor([
[0, 1],
[1, 2],
[1, 3],
])
L[i[:,:,None], i[:,None]]
輸出：
tensor([[[ 0, 1],
[ 4, 5]],
[[ 5, 6],
[ 9, 10]],
[[ 5, 7],
[13, 15]]])
總而言之，我建議使用索引而不是掩碼。
轉載請註明出處，本文鏈接：https://www.uj5u.com/caozuo/441538.html
上一篇：拆分Csv檔案中的列