PyTorch torch.vmap 関数
Pytorch torch リファレンスマニュアルtorch.vmapPyTorchにおけるベクトルマッピング用の関数です。関数を入力として受け取り、バッチのテンソルに対して自動的に操作できる新しい関数を返します。JAXのvmapに似ています。
関数定義
torch.vmap(func, in_dims, out_dims, randomness)
パラメータ説明
func: ベクトル化する関数in_dims: 入力テンソルのバッチ次元（オプション）out_dims: 出力テンソルのバッチ次元（オプション）randomness: ランダム動作。"error"、"different"、"same"から選択可能
使用例
例
import torch
# 簡単な関数を定義する
def simple_func(x):
return x * 2 + 1
# vmapを使用して関数をベクトル化する
vectorized_func = torch.vmap(simple_func)
# バッチ入力（バッチ次元は0）
batch_input = torch.randn(4, 3)
# ベクトル化された関数を適用する
output = vectorized_func(batch_input)
print("入力形状:", batch_input.shape)
print("出力形状:", output.shape)
print("出力:")
print(output)
# 簡単な関数を定義する
def simple_func(x):
return x * 2 + 1
# vmapを使用して関数をベクトル化する
vectorized_func = torch.vmap(simple_func)
# バッチ入力（バッチ次元は0）
batch_input = torch.randn(4, 3)
# ベクトル化された関数を適用する
output = vectorized_func(batch_input)
print("入力形状:", batch_input.shape)
print("出力形状:", output.shape)
print("出力:")
print(output)
出力結果は次のとおりです：
输入形状: torch.Size([4, 3])
输出形状: torch.Size([4, 3])
输出:
tensor([[ 0.2345, 1.5678, -0.3456],
[ 2.1234, -1.2345, 0.5678],
[-0.8765, 1.2345, 2.3456],
[ 1.5678, 0.1234, -1.2345]])
その他の拡張