在 Python 裡面用索引取陣列裡面的值，通常我們會這麼寫：
var = tensor[index]
進階一點，如果你今天得到的索引值是一張陣列，而且你需要得到一張陣列的返回值呢？
在 C 語言裡面，比較直覺的方法可能是 For loop，
但是我們都用到 Python 了，該用高級一點的方法了！
以下介紹 Python （list, numpy 適用）中 PyTorch 的 tensor 陣列索引賦值法：
data = np.random.randint(0,1023, size=(32, 32)) data = torch.from_numpy(data) index = np.argmin(data, axis=1) data = data.contiguous().view(-1) base = np.arange(0,32)*32 index = index + base output = data[index]
上面程式碼等效：
output = np.min(data, axis=1)
這樣你會問我：哇！Ben 為什麼一行能解決的問題，你要這麼費工？
傻傻的，
今天是因為有 np.min() 可以用，如果你今天要取得的並不是這麼簡單的統計值呢？
這種賦值法的價值在於核心演算法複雜時，
還能提供高效率的計算複雜度。
讓我們來看一個比較實際的例子，
當你今天是一個通訊所的學生，老師要你量化 16-QAM 的訊號：
B, C, H = 32000, 3, 16 Q = 16 data = np.random.randint(0,1023, size=(B, C, H)) data = torch.from_numpy(data) data = data.contiguous().view(-1) ################### Kernel ######################### q_list = torch.unsqueeze(torch.arange(1,Q+1)*64, dim=0) data = torch.unsqueeze(data, dim=1).repeat(1, Q) q_matrix = q_list.repeat(int(list(data.size())[0]), 1) index = np.argmin((data - q_matrix)**2, axis=1) feature = q_list[0, index] #########################################################
上述核心來自於我先前提到的 KNN 算法。
既然都算出來了，讓我們來比較與 double for loop 賦值 的速度：
f1s = time.time()
feature = q_list[0, index]
f1e = time.time()
print("non-loop processing time: " + "%.3f" %(f1e-f1s))
feature2 = torch.ones(int(list(data.size())[0]))
f2s = time.time()
k=0
for i in index:
for j in torch.arange(0,len(q_list[0]+1)):
if j == i:
feature2[k] = q_list[0][int(j)]
k=k+1
break
f2e = time.time()
print("for-loop processing time: " + "%.3f" %(f2e-f2s))
feature = feature.type(torch.FloatTensor)
print("feature = feature2? : " + str(torch.equal(feature, feature2)))
 |
| 差了六十萬倍呢！ |
結論：當數據足夠大時，陣列索引賦值法會與迴圈賦值法速度差距會越來越明顯。
如果對於文章內容有疑問的，歡迎聯絡我： wuyiulin@gmail.com