Функция PyTorch torch.set_num_threads
Справочное руководство PyTorch torch
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])
Другие расширения