Función PyTorch torch.set_float32_matmul_precision

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

torch.set_float32_matmul_precisionEs una función en PyTorch utilizada para establecer la precisión de la multiplicación de matrices float32. Se puede optar por usar lower precision para mejorar el rendimiento, o higher precision para mejorar la precisión.

Definición de la función

torch.set_float32_matmul_precision(precision)

Descripción de parámetros

  • precision: Nivel de precisión, valores opcionales:
    • "highest": Máxima precisión (por defecto)
    • "high": Alta precisión
    • "medium": Precisión media (usa TensorFloat-32)

Ejemplo de uso

Ejemplo

import torch

# Establecer precisión media (aceleración con TensorFloat-32)
torch.set_float32_matmul_precision("medium")

# Crear matrices para probar
a = torch.randn(100, 100)
b = torch.randn(100, 100)

# Multiplicación de matrices
c = torch.matmul(a, b)

print("Realizar multiplicación de matrices con precisión media")
print("Forma del resultado:", c.shape)

# Restaurar la precisión máxima
torch.set_float32_matmul_precision("highest")

El resultado de salida es:

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

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

Otras extensiones