スポンサーリンク

Tensor配列の最大値と最小値を同時に取得するtorch.aminmax【PyTorch】

Python
2022/09/22

PyTorchのTensor配列の最大値と最小値を同時に取得するにはtorch.aminmaxを使用する。

torch.aminmaxの使い方

torch.aminmaxにTensor配列を指定することで、最大値と最小値を取得することができる。また、最大値を得たい場合はmaxメソッド、最小値を得たい場合はminメソッドを用いる。Tensor配列にnanが含まれている場合は、nanが取得される。

import torch

a = torch.tensor([1, -3, 5])

print(torch.aminmax(a))
# torch.return_types.aminmax(
# min=tensor(-3),
# max=tensor(5))

print(torch.aminmax(a).max)
# tensor(5)
print(torch.aminmax(a).min)
# tensor(-3)

b = torch.tensor([1, -3, 5, torch.nan])
print(torch.aminmax(b))
# torch.return_types.aminmax(
# min=tensor(nan),
# max=tensor(nan))

多次元配列の場合、引数dimを指定すると、指定した軸に沿った最大値と最小値が取得される。

c = torch.tensor([[0, 1, 2],
                  [3, 4, 5]])
print(torch.aminmax(c))
# torch.return_types.aminmax(
# min=tensor(0),
# max=tensor(5))

print(torch.aminmax(c, dim=0))
# torch.return_types.aminmax(
# min=tensor([0, 1, 2]),
# max=tensor([3, 4, 5]))

print(torch.aminmax(c, dim=1))
# torch.return_types.aminmax(
# min=tensor([0, 3]),
# max=tensor([2, 5]))

関連記事、参照記事

最大値や最小値を個別に取得したい場合は、torch.maxやtorch.minを使用します。

コメント