Функция PyTorch torch.set_float32_matmul_precision

Pytorch torch 参考手册Справочное руководство Pytorch torch

torch.set_float32_matmul_precision— это функция в PyTorch, используемая для установки точности умножения матриц float32. Вы можете выбрать использование более низкой точности для повышения производительности или более высокой точности для повышения точности вычислений.

Определение функции

torch.set_float32_matmul_precision(precision)

Описание параметра

  • precision: Уровень точности, возможные значения:
    • "highest": Наивысшая точность (по умолчанию)
    • "high": Высокая точность
    • "medium": Средняя точность (использует TensorFloat-32)

Пример использования

Пример

import torch

# Установить среднюю точность (использовать TensorFloat-32 для ускорения)
torch.set_float32_matmul_precision("medium")

# Создать матрицы для тестирования
a = torch.randn(100, 100)
b = torch.randn(100, 100)

# Умножение матриц
c = torch.matmul(a, b)

print("Умножение матриц с использованием средней точности")
print("Форма результата:", c.shape)

# Восстановить наивысшую точность
torch.set_float32_matmul_precision("highest")

Вывод:

使用中等精度进行矩阵乘法
结果形状: torch.Size([100, 100])

Pytorch torch 参考手册Справочное руководство Pytorch torch

Другие расширения