PyTorchでテンソルの平均と標準偏差を計算する方法｜torch.mean()・torch.std()の使い方

PyTorchのテンソルはNumPy配列によく似ています。最大の違いは、テンソルがGPUを活用して数値計算を高速化できる点です。テンソルの平均はtorch.mean()メソッドで計算でき、入力テンソル内の全要素の平均値が返されます。さらに、適切な軸（dim）を指定すれば、行ごと・列ごとの平均も求められます。

テンソルの標準偏差はtorch.std()メソッドで計算します。テンソル内の全要素に対する標準偏差を返し、平均と同様に行単位・列単位の計算も可能です。なお、torch.std()はデフォルトで不偏標準偏差（ベッセル補正あり）を返す点に注意してください。

計算の手順

• 必要なライブラリをインポートします。以降のPythonサンプルではtorchライブラリを使用するため、事前にインストールしておいてください。

• PyTorchテンソルを定義し、内容をprint()で確認します。

• torch.mean(input, dim)で平均を計算します。inputは平均を求める対象のテンソル、dim（axisでも可）は集計対象の次元です。計算結果は新しい変数に代入します。

• torch.std(input, dim)で標準偏差を計算します。引数の意味はtorch.mean()と同じです。計算結果は新しい変数に代入します。

• 計算した平均と標準偏差を出力して確認します。

例1：1次元テンソルの場合

次のPythonプログラムは、1次元テンソルの平均と標準偏差を計算する例です。

# Pythonプログラム：1次元テンソルの平均と標準偏差を計算
# ライブラリをインポート
import torch

# テンソルを作成
T = torch.Tensor([2.453, 4.432, 0.754, -6.554])
print("T:", T)

# 平均と標準偏差を計算
mean = torch.mean(T)
std = torch.std(T)

# 結果を出力
print("Mean:", mean)
print("Standard deviation:", std)

出力

T: tensor([ 2.4530,  4.4320,  0.7540, -6.5540])
Mean: tensor(0.2713)
Standard deviation: tensor(4.7920)

例2：2次元テンソルの場合（行・列方向の集計）

次のPythonプログラムは、2次元テンソルの平均と標準偏差を、テンソル全体・列方向・行方向のそれぞれについて計算する例です。

# 必要なライブラリをインポート
import torch

# 3x4の2次元テンソルを作成
T = torch.Tensor([[2,4,7,-6],
[7,33,-62,23],
[2,-6,-77,54]])
print("T:\n", T)

# 全体の平均と標準偏差を計算
mean = torch.mean(T)
std = torch.std(T)
print("Mean:", mean)
print("Standard deviation:", std)

# 列方向（axis=0）の平均と標準偏差を計算
mean = torch.mean(T, axis = 0)
std = torch.std(T, axis = 0)
print("Column-wise Mean:\n", mean)
print("Column-wise Standard deviation:\n", std)

# 行方向（axis=1）の平均と標準偏差を計算
mean = torch.mean(T, axis = 1)
std = torch.std(T, axis = 1)
print("Row-wise Mean:\n", mean)
print("Row-wise Standard deviation:\n", std)

出力

T:
tensor([[  2.,   4.,   7.,  -6.],
        [  7.,  33., -62.,  23.],
        [  2.,  -6., -77.,  54.]])
Mean: tensor(-1.5833)
Standard deviation: tensor(36.2703)
Column-wise Mean:
tensor([  3.6667,  10.3333, -44.0000,  23.6667])
Column-wise Standard deviation:
tensor([ 2.8868, 20.2567, 44.7996, 30.0056])
Row-wise Mean:
tensor([  1.7500,   0.2500,  -6.7500])
Row-wise Standard deviation:
tensor([ 5.5603, 42.8593, 53.8602])

まとめ

torch.mean()とtorch.std()を使えば、テンソル全体の統計量はもちろん、dim（axis）を指定するだけで行・列ごとの集計も簡単に行えます。なお、近年のPyTorchでは標準偏差計算時に補正方法を明示することが推奨されており、torch.std(T, correction=1)のように書くと意図がより明確になります。データ分析や前処理において、これらの統計量は正規化やスケーリングの基礎となるため、ぜひ使いこなしましょう。