PyTorch torch.set_float32_matmul_precision 関数

Pytorch torch 参考手册Pytorch torch リファレンスマニュアル

torch.set_float32_matmul_precisionPyTorch で float32 行列乗算の精度を設定する関数です。lower precision を使用してパフォーマンスを向上させるか、higher precision を使用して正確性を向上させるかを選択できます。

関数定義

torch.set_float32_matmul_precision(precision)

パラメータ説明

  • precision: 精度レベル、指定可能な値：
    • "highest": 最高精度（デフォルト）
    • "high": 高精度
    • "medium": 中精度（TensorFloat-32 を使用）

使用例

例

import torch

# 中精度に設定（TensorFloat-32 を使用して高速化）
torch.set_float32_matmul_precision("medium")

# テスト用の行列を作成
a = torch.randn(100, 100)
b = torch.randn(100, 100)

# 行列乗算
c = torch.matmul(a, b)

print("中精度で行列乗算を実行")
print("結果の形状:", c.shape)

# 最高精度に戻す
torch.set_float32_matmul_precision("highest")

出力結果は：

使用中等精度进行矩阵乘法
结果形状: torch.Size([100, 100])

Pytorch torch 参考手册Pytorch torch リファレンスマニュアル

その他の拡張