PyTorch torch.set_num_threads 関数
torch.set_num_threadsこれは PyTorch で内部スレッド数を設定するための関数です。これらのスレッドは、テンソル操作(行列乗算、畳み込みなど)を並列実行するために使用されます。
関数定義
torch.set_num_threads(num_threads)
パラメータ説明
num_threads: 設定するスレッド数
使用例
例
import torch
# スレッド数を設定
torch.set_num_threads(4)
print("スレッド数は次のように設定されました:", torch.get_num_threads())
# テスト用に大きな行列を作成
a = torch.randn(1000, 1000)
b = torch.randn(1000, 1000)
# 行列乗算(設定されたスレッド数が使用されます)
c = torch.matmul(a, b)
print(行列乗算が完了しました。結果の形状:, c.shape)
# デフォルト値に戻す
torch.set_num_threads(1)
# スレッド数を設定
torch.set_num_threads(4)
print("スレッド数は次のように設定されました:", torch.get_num_threads())
# テスト用に大きな行列を作成
a = torch.randn(1000, 1000)
b = torch.randn(1000, 1000)
# 行列乗算(設定されたスレッド数が使用されます)
c = torch.matmul(a, b)
print(行列乗算が完了しました。結果の形状:, c.shape)
# デフォルト値に戻す
torch.set_num_threads(1)
出力結果は次のとおりです:
线程数已设置为: 4 矩阵乘法完成,结果形状: torch.Size([1000, 1000])
その他の拡張
Pytorch torch リファレンスマニュアル