問題場景描述

我在復(fù)現(xiàn)Faster-RCNN模型的過程中遇到這樣一個(gè)問題：
有一個(gè)張量，它的形狀是 (128, 21, 4)

roi_loc.shape = (128, 21, 4)

與之對(duì)應(yīng)的還有一個(gè)label數(shù)據(jù)

gt_label.shape = (128)

我現(xiàn)在的需求是將label當(dāng)作第一個(gè)張量在dim=1上的索引，將其中的數(shù)據(jù)拿出來。
具體來說就是，現(xiàn)在有128個(gè)樣本數(shù)據(jù)，每個(gè)樣本中有21個(gè)長度為4的向量。label也是128個(gè)，每個(gè)值代表取出21個(gè)向量中的哪一個(gè)。

問題的思考

我嘗試了很多辦法，包括布爾索引，index_select方法等，最后發(fā)現(xiàn)都不適用（也有可能我沒用好）。最后利用gather API解決了這個(gè)問題。

這個(gè)API的說明我看了很多遍都沒看懂，我相信絕大部分讀者也是因?yàn)榭床欢@個(gè)說明才來這兒的。
下面我給出自己的一些理解：

gather的說明

gather所需要的第一個(gè)參數(shù)是待索引的數(shù)據(jù)，在我們的問題中 roi_loc就是這個(gè)input。第二個(gè)參數(shù)dim，是你的索引數(shù)據(jù)要作用在哪個(gè)軸上，正如前面所言，我們想索引第二個(gè)軸(dim=1).
最難理解的是index，index就是我們想要用來索引的張量，對(duì)應(yīng)的是label?？墒莑abel不能直接拿來用，得先做一定的變換，這也就是gather的難點(diǎn)。

我們先從簡單的情況來看

input和gather必須在維度上相同，假設(shè)數(shù)據(jù)還是3 * 3，index也是1 * 3的(注意這里是二維的)

此時(shí)row至多取值0，col至多取值為2
如果我要對(duì)dim=0索引
那么data[0][0] = data[index[0][0]] [0] = data[1][0] = 2
data[0][1] = data[index[0][1]] [1] = data[0][1] = 5
data[0][2] = data[index[0][2]][2] = data[2][2] = 9
上面的過程可以描述為，第一列的元素我想選第二行的，第二列的元素我想選第一行的，第三列的元素我想選第三行的。
可以發(fā)現(xiàn)因?yàn)閕ndex是1 * 3的，所以最后的輸出也是31* 3，即輸出張量的shape取決于index的shape

以上過程我相信讀者好好體悟應(yīng)該可以理解。

問題的解決

回到我們的問題

roi_loc.shape = (128, 21, 4)，gt_label.shape = (128)

我們想索引dim=1，最后的結(jié)果應(yīng)該是(128, 4)
由上面的說明可以知道，input和index的dimension首先得相同

idx = gt_roi_labels.unsqueeze(-1).unsqueeze(-1)
idx.shape = (128, 1, 1)

又因?yàn)槲覀兿胍敵龅慕Y(jié)果得是（128， 4），所以得讓idx在最后一個(gè)軸上重復(fù)4次

idx = idx.repeat_interleave(-1, 4)
idx.shape = (128, 1, 4)

現(xiàn)在就可以利用gather在dim=1上索引了

result = roi_loc.gather(1, idx)
result.shape = (128, 1, 4)

最后將長度為1的軸壓縮(本身這個(gè)軸的出現(xiàn)是為了滿足input和index維度一樣的要求)

result = result.squeeze(1)
result.shape(128, 4)

以上就是Pytorch深度學(xué)習(xí)gather一些使用問題解決方案的詳細(xì)內(nèi)容，更多關(guān)于Pytorch學(xué)習(xí)gather使用問題的資料請(qǐng)關(guān)注腳本之家其它相關(guān)文章！

最新評(píng)論