Fonction torch.set_num_threads de PyTorch

Pytorch torch 参考手册Manuel de référence PyTorch torch

torch.set_num_threadsest une fonction dans PyTorch utilisée pour définir le nombre de threads internes. Ces threads sont utilisés pour exécuter des opérations tensorielles en parallèle (comme la multiplication matricielle, la convolution, etc.)

Définition de fonction

torch.set_num_threads(num_threads)

Description du paramètre

  • num_threads: Le nombre de threads à définir

Exemple d'utilisation

Exemple

import torch

# Définir le nombre de threads
torch.set_num_threads(4)

print("Nombre de threads défini à :", torch.get_num_threads())

# Créer de grandes matrices pour les tests
a = torch.randn(1000, 1000)
b = torch.randn(1000, 1000)

# Multiplication matricielle (utilisera le nombre de threads défini)
c = torch.matmul(a, b)

print("Multiplication matricielle terminée, forme du résultat :", c.shape)

# Restaurer les valeurs par défaut
torch.set_num_threads(1)

Le résultat de sortie est:

线程数已设置为: 4
矩阵乘法完成，结果形状: torch.Size([1000, 1000])

Pytorch torch 参考手册Manuel de référence PyTorch torch

Autres extensions