在閱讀使用 pytorch 實現的代碼時，筆者會遇到需要對某一維數據進行求和( sum )或 softmax 的操作。在 pytorch 中，上述兩個方法均帶有一個指定維度的 dim 參數，這里記錄下 dim 參數的用法。
torch.sum
在 pytorch 中，提供 torch.sum 的兩種形式，一種直接將待求和數據作為參數，則返回參數數據所有維度所有元素的和，另外一種除接收待求和數據作為參數外，還可加入 dim 參數，指定對待求和數據的某一維進行求和。
out = torch.sum( a ) #對 a 中所有元素求和 out = torch.sum( a , dim = 1 ) #對 a 中第 1 維的元素求和
上述第一種形式比較好理解，但第二種形式，加入 dim 參數后，比較令人疑惑的是到底哪些元素參與了求和？這里通過例子來進行說明。
1）首先我們生成一個維度為 ( 3, 4, 5, 6 ) 的元素全為 1.0 的 tensor a。
>>> import torch >>> a = torch.ones( 3, 4, 5, 6 ) #生成一個形狀為 ( 3, 4, 5, 6 ) 的數據,數據類型默認為 torch.FloatTensor
2）使用 sum 方法對上述生成的 tensor 進行求和操作。注意 tensor 的維度索引從 0 開始。
>>> b = torch.sum( a ) #對 a 中所有元素求和， b = 360.0 >>> c = torch.sum( a, dim = 0 ) #對 a 中 dim = 0 元素求和 >>> c.shape # c 的 shape 為 torch.Size( [ 4, 5, 6 ] )，其中所有元素值為 3.0 >>> d = torch.sum( a, dim = 3 ) #對 a 中 dim = 3 元素求和 >>> d.shape # d 的 shape 為 torch.Size( [ 3, 4, 5 ] ),其中所有元素值為 6.0
對上述結果進行解釋，b 的結果很好理解，因為 tensor a 的維度為 ( 3, 4, 5, 6 ) 且其中所有元素的值為 1，則對其中所有元素求和的結果為 3 * 4 * 5 * 6 * 1.0 = 360.0 .
對於 c 和 d 的結果，首先可以觀察得到的是， 若在第 i 維進行求和，即 sum( a, dim = i )，則求和結果的每一個元素的值均為該維度的大小。如在 dim = 0 求和，在 dim = 0 上 a 的尺寸為 3，則求和結果 c 的每一個元素值為 3.0 .也就是說每個結果元素值均為是三個求和元素值( 1.0 )相加的結果，求和結果 c 的維度為 ( 4, 5, 6 )，說明待求和數據 a 分為 ( 4, 5, 6 ) 共 4 * 5 * 6 組的元素進行了求和運算。在 dim = 3 上的求和結果 d 現象與 c 保持一致。
對於輸入待求和數據所有數據元素均為 1 時，可以歸納出一個結論，對於維度為 ( s0, s1, s2, s3 ) 的 tensor 的第 i 維進行求和，如第 0 維，則結果的維度為 ( s1, s2, s3 )，其維度為原輸入維度去除求和維度。結果的每一個元素值即為 1 * s0 = s0，即為待求和維度的尺寸。
下面以三維數據即維度為 ( 3, 4, 4 ) 的 tensor a 為例展示 sum 在某一維度的實際計算過程。
使用 dim = 0 參數計算時，產生的結果維度為 ( 4, 4 ), 對於結果中的每一個位置 ( i, j ) ，由 3 個元素進行計算，實際計算的是 a[ 0 ][ i ][ j ] + a[ 1 ][ i ][ j ] + a[ 2 ][ i ][ j ]，當上述三個元素的值均為 1.0 時，計算結果元素即為 3.0 。如上圖左側的圖，a[ 0 ][ 3 ][ 3 ] + a[ 1 ][ 3 ][ 3 ] + a[ 2 ][ 3 ][ 3 ] 的結果即為輸出 ( 3, 3 ) 位置上的值。上述位置索引 ( i, j ) 的數量由輸入的待求和數據的其他維度的尺寸決定。
使用 dim = 2 參數計算時，產生的結果維度為 ( 3, 4 ),對於結果中的每一個位置( i, j ) ，由 4 個元素進行計算，實際計算的是 a[ i ][ j ][ 0 ] + a[ i ][ j ][ 1 ] + a[ i ][ j ][ 2 ] + a[ i ][ j ][ 3 ]，當上述四個元素的值均為 1.0 時，計算結果元素即為 4.0 。如 a[ 0 ][ 0 ][ 0 ] + a[ 0 ][ 0 ][ 1 ] + a[ 0 ][ 0 ][ 2 ] + a[ 0 ][ 0 ][ 3 ] 即為輸出 ( 0, 0 ) 位置上的值。
對於維度為 ( s0, s1, s2, ... , si, ... , sn ) 的待求和向量，使用 dim = i 調用 sum 方法，則實際產生的結果維度為 ( s0, s1, s2, ... , si-1, si+1, ... , sn )，每個結果元素由 si 個元素元素求和獲得。這 si 個元素坐標在其他維度索引保持一致，而在待求和維度索引由 0 至 si 變化。可以看到共有 ( s0, s1, s2, ... , si-1, si+1, ... , sn ) 組這樣的求和元素( 索引的數量 )，即為結果的維度。
torch.nn.softmax / torch.nn.functional.softmax
softmax 是神經網路中常見的一種計算函數，其可將所有參與計算的對象值映射到 0 到 1 之間，並使得計算對象的和為 1. 在 pytorch 中的 softmax 方法在使用時也需要通過 dim 方法來指定具體進行 softmax 計算的維度。這里以 torch.nn.functional.softmax 為例進行說明。
softmax 在 pytorch 官方文檔中的描述如下：
It is applied to all slices along dim, and will re-scale them so that the elements lie in the range [0, 1] and sum to 1.
可以明確的是， softmax 計算獲得的數值在 0 - 1 之間，但是同樣比較令人疑惑的是，all slices along the dim 具體指代的是那些數據。這里使用一個維度為 ( 2, 2, 2 ) 的 tensor a 作為示例。
>>> import torch >>> import torch.nn.functional as f >>> a = torch.ones( 2, 2, 2 ) >>> b = f.softmax( a, dim=0 ) #對 a 的第 0 維進行 softmax 計算
與 sum 方法不同，softmax 方法計算獲得的結果的維度與輸入的待計算的數據的維度保持一致( sum 方法求和后進行指定求和的那一維不會出現在結果維度中 )。若對第 n 維進行 softmax 操作，且該維尺寸為 s，則 softmax 的結果為其他維索引保持一致，當前維索引由 0 至 s - 1 共 s 個值得和為 1.
參與 softmax 計算的元素與 sum 方法很相似，對於 tensor a 在 dim = 0 進行 softmax，輸出結果 b 實際上是 b[ 0 ][ i ][ j ] + b[ 1 ][ i ][ j ] 的值為 1.即其他維度索引保持一致，而在進行 softmax 維度索引由 0 至 si 變化，如 b[ 0 ][ 0 ][ 1 ] + a[ 1 ][ 0 ][ 1 ] 的值為1.對於 tensor a 在 dim = 2 進行 softmax，輸出結果 b 實際上是 b[ i ][ j ][ 0 ] + a[ i ][ j ][ 1 ] 的值為 1.
>>> c = b[ 0 ][ 0 ][ 1 ] + b[ 1 ][ 0 ][ 1 ] #c 的值為 1
參考