PyTorch torch.set_num_threads Function
PyTorch torch Reference Manual
torch.set_num_threadsis a function in PyTorch used to set the number of internal threads. These threads are used to execute tensor operations in parallel (such as matrix multiplication, convolution, etc.)
Function Definition
torch.set_num_threads(num_threads)
Parameter Description
num_threads: The number of threads to set
Usage Example
Example
import torch
# Set number of threads
torch.set_num_threads(4)
print("Number of threads set to:", torch.get_num_threads())
# Create large matrices for testing
a = torch.randn(1000, 1000)
b = torch.randn(1000, 1000)
# Matrix multiplication (will use the set number of threads)
c = torch.matmul(a, b)
print("Matrix multiplication completed, result shape:", c.shape)
# Restore default values
torch.set_num_threads(1)
# Set number of threads
torch.set_num_threads(4)
print("Number of threads set to:", torch.get_num_threads())
# Create large matrices for testing
a = torch.randn(1000, 1000)
b = torch.randn(1000, 1000)
# Matrix multiplication (will use the set number of threads)
c = torch.matmul(a, b)
print("Matrix multiplication completed, result shape:", c.shape)
# Restore default values
torch.set_num_threads(1)
The output result is:
线程数已设置为: 4 矩阵乘法完成，结果形状: torch.Size([1000, 1000])
Other Extensions