Función torch.set_num_threads de PyTorch

Pytorch torch 参考手册Manual de referencia de Pytorch torch

torch.set_num_threadsEs una función de PyTorch que se utiliza para establecer el número de hilos internos. Estos hilos se utilizan para ejecutar operaciones con tensores en paralelo (como multiplicación de matrices, convoluciones, etc.).

Definición de la función

torch.set_num_threads(num_threads)

Descripción de parámetros

  • num_threads: El número de hilos a establecer

Ejemplo de uso

Ejemplo

import torch

# Establecer el número de hilos
torch.set_num_threads(4)

print("El número de hilos se ha establecido en:", torch.get_num_threads())

# Crear una matriz grande para probar
a = torch.randn(1000, 1000)
b = torch.randn(1000, 1000)

# Multiplicación de matrices (usará el número de hilos configurado)
c = torch.matmul(a, b)

print("Multiplicación de matrices completada, forma del resultado:", c.shape)

# Restaurar el valor predeterminado
torch.set_num_threads(1)

El resultado de salida es:

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

Pytorch torch 参考手册Manual de referencia de Pytorch torch

Otras extensiones