Pythonで引数などによって各methodを実行する方法

Question

やりたいこと。

import torch
torch.nn.MSELoss()

このMSELoss()というメソッドをユーザーに選択させたい場合、古典的な方法では

# argpase は使わない
method = "MSE" # JupyterNotebookなどでユーザーに直接指定させます

# ここから古典的な手法
if method=="MSE":
    torch.nn.MSELoss()
elif method=="BCE":
    torch.nn.BCELoss()
else:
    torch.nn.CrossEntropyLoss

となります。これをif文を使わずに、変数自体でメソッドを実行できる方法があれば教えて下さい。或いは「そのような方法は存在しない」でお願いすします。

Answer 1

以下でできます。

方法1

名前 → クラス の dict を作っておく

import torch

name = "MSE"

losses = {
    "MSE": torch.nn.MSELoss,
    "BCE": torch.nn.BCELoss,
    "CrossEntropy": torch.nn.CrossEntropyLoss,
}

try:
    loss = losses[name]
except:
    raise ValueError("存在しません")
    
print(loss)
# <class 'torch.nn.modules.loss.MSELoss'>
# loss() で呼び出せる

dict を自動で作るなら

import torch
from pprint import pprint

def get_losses():
    return {x: getattr(torch.nn, x) for x in dir(torch.nn) if x.endswith("Loss")}

losses = get_losses()

pprint(losses)

方法2

getattr() で torch.nn 以下にその名前のクラスがあれば取得する

import torch

name = "MSE"

def get_losses():
    return "\n".join([x for x in dir(torch.nn) if x.endswith("Loss")])

try:
    loss = getattr(torch.nn, name + "Loss")
except:
    raise ValueError(f"{name}Loss は存在しません。利用可能な Loss 一覧\n{get_losses()}")
    
print(loss)
# <class 'torch.nn.modules.loss.MSELoss'>
# loss() で呼び出せる

Comments

[teefpc]

なるほど。これは凄くいいですね。
どうもありがとうございました！