【Pytorch】torch.argmaxの使い方・引数を徹底解説！インデックスを取得する理由についても解説！

import torch

# テンソルの作成
tensor = torch.tensor([1, 3, 5, 2, 4])

# 最大値のインデックスを取得
max_index = torch.argmax(tensor)

print(max_index)  # 出力: tensor(2)

# テンソルの特定のインデックスの値を取得
max_value = tensor[max_index]

print(max_value)  # 出力: tensor(5)

import torch

# 1次元テンソルの場合
input_tensor = torch.tensor([3, 1, 5, 2, 4])
max_index = torch.argmax(input_tensor)
print(max_index)  # 出力: tensor(2)
# input_tensorの最大値のインデックスは2

# 2次元テンソルの場合
input_tensor = torch.tensor([[3, 1, 5], [2, 4, 6]])
max_index = torch.argmax(input_tensor, dim=1)
print(max_index)  # 出力: tensor([2, 2])
# dim=1により、各行で最大値のインデックスを取得
# 最初の行の最大値のインデックスは2、2番目の行の最大値のインデックスも2

# 2次元テンソルの場合
input_tensor = torch.tensor([[3, 1, 5], [2, 4, 6]])
max_index = torch.argmax(input_tensor, dim=-1)
print(max_index)  # 出力: tensor([2, 2])
# dim=-1は最後の次元を表し、この場合は各行での最大値のインデックスを取得しています

# 3次元テンソルの場合
input_tensor = torch.tensor([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]])
max_index = torch.argmax(input_tensor, dim=-2)
print(max_index)  # 出力: tensor([[1, 1, 1], [1, 1, 1]])
# dim=-2は最後から2番目の次元を表し、この場合は各列ごとに最大値のインデックスを取得しています

import torch

# 2次元テンソルの場合
input_tensor = torch.tensor([[3, 1, 5], [2, 4, 6]])
max_index = torch.argmax(input_tensor, dim=1, keepdim=False)
print(max_index)  # 出力: tensor([2, 2])
# keepdim=Falseの場合、出力テンソルは指定した次元を保持せずに返される

max_index = torch.argmax(input_tensor, dim=1, keepdim=True)
print(max_index)  # 出力: tensor([[2], [2]])
# keepdim=Trueの場合、出力テンソルは指定した次元を保持して返される

# 3次元テンソルの場合
input_tensor = torch.tensor([[[3, 1, 5], [2, 4, 6]], [[7, 2, 4], [6, 9, 2]]])
max_index = torch.argmax(input_tensor, dim=2, keepdim=False)
print(max_index)
# 出力: tensor([[2, 2],
#                [0, 1]])
# keepdim=Falseの場合、出力テンソルは指定した次元を保持せずに返される

max_index = torch.argmax(input_tensor, dim=2, keepdim=True)
print(max_index)
# 出力: tensor([[[2],
#                [2]],
#
#               [[0],
#                [1]]])
# keepdim=Trueの場合、出力テンソルは指定した次元を保持して返される

import torch

# 1次元テンソルの場合
input_tensor = torch.tensor([3, 1, 5, 5, 2, 4])
max_index = torch.argmax(input_tensor)
print(max_index)  # 出力: tensor(2)
# input_tensorの最大値のインデックスは2

# 複数の最大値がある場合
input_tensor = torch.tensor([3, 1, 5, 5, 2, 4])
max_indices = torch.where(input_tensor == torch.max(input_tensor))[0]
print(max_indices)  # 出力: tensor([2, 3])
# torch.whereを使用して、最大値と同じ値を持つインデックスを取得

# インデックスをランダムに選択する場合
import random
random_max_index = random.choice(max_indices)
print(random_max_index)  # 出力: 2または3
# max_indicesからランダムに1つのインデックスを選択する

# すべての最大値のインデックスを取得する場合
all_max_indices = torch.nonzero(input_tensor == torch.max(input_tensor)).squeeze()
print(all_max_indices)  # 出力: tensor([2, 3])
# torch.nonzeroを使用して、最大値と同じ値を持つすべてのインデックスを取得
# squeeze()を使用して、不要な次元を削除