スポンサーリンク

【PyTorch】平均値と中央値を算出するtorch.mean、torch.median

Python

PyTorchのTensor配列で平均値を求めるにはtorch.meantorch.nanmean、中央値を求めるにはtorch.mediantorch.nanmedianを使う。

平均値:torch.mean

torch.mean()の第一引数inputにTensor配列を渡すと、平均値が返ってくる。

import torch

a = torch.arange(12, dtype=torch.float32).reshape(3, 4)
print(a.shape)
print(a)
# torch.Size([3, 4])
# tensor([[ 0,  1,  2,  3],
#        [ 4,  5,  6,  7],
#        [ 8,  9, 10, 11]])

print(torch.mean(input=a))
# tensor(5.5000)

引数dimに0を指定するて列ごとの平均値、dimに1を指定すると行ごとの平均値を算出できる。

print(torch.mean(input=a, dim=0))
# tensor([4., 5., 6., 7.])

print(torch.mean(input=a, dim=1))
# tensor([1.5000, 5.5000, 9.5000])

nanがあるTensor配列をtorch.meanに渡すと、nanが返ってくる。

b = torch.tensor([[torch.nan, 1, 2], [1, 2, 3]])

print(b.shape)
print(b)
# torch.Size([2, 3])
# tensor([[nan, 1., 2.],
#        [1., 2., 3.]])

print(torch.mean(input=b))
# tensor(nan)

torch.nanmean

nanを除外して平均値を求めたい場合は、torch.nanmeanを使う。使い方は、torch.meanと同様である。

print(torch.nanmean(input=b))
# tensor(1.8000)

print(torch.nanmean(input=b, dim=0))
# tensor([1.0000, 1.5000, 2.5000])

print(torch.nanmean(input=b, dim=1))
# tensor([1.5000, 2.0000])

中央値:torch.median

torch.median()の第一引数inputにTensor配列を渡すと、中央値が返ってくる。

print(torch.median(input=a))
# tensor(5.)

引数dimに0を指定するて列ごとの中央値、dimに1を指定すると行ごとの中央値を算出できる。また中央値のインデックスが返ってくる。

values, indices = torch.median(input=a, dim=0)
print(values)
print(indices)
# tensor([4., 5., 6., 7.])
# tensor([1, 1, 1, 1])

values, indices = torch.median(input=a, dim=1)
print(values)
print(indices)
# tensor([1., 5., 9.])
# tensor([1, 1, 1])

nanがあるTensor配列をtorch.medianに渡すと、nanが返ってくる。

print(torch.median(input=b))
# tensor(nan)

torch.nanmedian

nanを除外して中央値を求めたい場合は、torch.nanmedianを使う。使い方は、torch.medianと同様である。

print(torch.nanmedian(input=b))
# tensor(2.)

values, indices = torch.nanmedian(input=b, dim=0)
print(values)
print(indices)
# tensor([1., 1., 2.])
# tensor([1, 0, 0])

values, indices = torch.nanmedian(input=b, dim=1)
print(values)
print(indices)
# tensor([1, 0, 0])
# tensor([1, 1])

関連記事、参考資料

PyTorch公式ページでも紹介されていた本で、「Tensorの仕組み」から「ディープラーニングの実践プロジェクト:肺がんの早期発見」までステップバイステップで説明しているため、PyTorchの中身をよく理解できるのでオススメです。

PyTorch実践入門 (Compass Booksシリーズ)

コメント