NumPy/TFv1转PyTorch后函数结果精度损失问题咨询

问题描述 #

尝试将模型从TensorFlow v1迁移至PyTorch时，发现某导数计算函数的PyTorch版本返回结果与NumPy/TensorFlow v1版本存在约10%的偏差，该偏差导致损失函数异常，模型无法训练。已确认：

• PyTorch模型的权重与偏差和NumPy版本完全一致
• NumPy结果与TensorFlow v1结果精度吻合度达1e-5，原模型可正常训练

现需解答：

1. 为何NumPy函数结果精度更优？
2. 如何调整PyTorch函数使其精度接近NumPy版本？

数据加载及函数实现代码 #

import pickle
import numpy as np
import torch
from typing import Dict, Any

with open('recovered_autoencoder_network.pkl', 'rb') as f:
    recovered_autoencoder_network = pickle.load(f)

# parameters needed for this issue
params: Dict[str, Any] = {'weight_precision': torch.float64,
                          'sindy_precision': torch.float64,
                          'target_device': 'cuda'}

sindy_autoencoder = torch.load('saved_model.pkl')
sindy_autoencoder.to(params['target_device'])

# this is a version of the 'problem' function in torch.
def calculate_first_and_second_derivative_with_torch(input_and_derivatives, stack):

    x, dx, ddx = input_and_derivatives

    layer_count = len(stack)

    for i in range(layer_count - 1):
        x = torch.mm(x, stack[i].weights) + stack[i].bias
        x = torch.sigmoid(x)
        dx_prev = torch.mm(dx, stack[i].weights)
        sigmoid_first_derivative = torch.mul(x, 1 - x)
        sigmoid_second_derivative = torch.mul(sigmoid_first_derivative, 1 - 2 * x)
        dx = torch.mul(sigmoid_first_derivative, dx_prev)
        ddx = torch.mul(sigmoid_second_derivative, torch.square(dx_prev)) \
              + torch.mul(sigmoid_first_derivative, torch.mm(ddx, stack[i].weights))
    dx = torch.mm(dx, stack[layer_count - 1].weights)
    ddx = torch.mm(ddx, stack[layer_count - 1].weights)

    return dx, ddx

# this is the equivalent 'problem' function in numpy.
def calculate_first_and_second_derivative_with_np(input, dx, ddx, weights, biases):
    dz = dx
    ddz = ddx

    def sigmoid(x):
        return 1 / (1 + np.exp(-x))

    for i in range(len(weights) - 1):
        input = np.matmul(input, weights[i]) + biases[i]
        input = sigmoid(input)
        dz_prev = np.matmul(dz, weights[i])
        sigmoid_derivative = np.multiply(input, 1 - input)
        sigmoid_derivative2 = np.multiply(sigmoid_derivative, 1 - 2 * input)
        dz = np.multiply(sigmoid_derivative, dz_prev)
        ddz = np.multiply(sigmoid_derivative2, np.square(dz_prev)) \
              + np.multiply(sigmoid_derivative, np.matmul(ddz, weights[i]))
    dz = np.matmul(dz, weights[-1])
    ddz = np.matmul(ddz, weights[-1])

    return dz, ddz

dx_decode_np_test, ddx_decode_np_test = \
    calculate_first_and_second_derivative_with_np(
        recovered_autoencoder_network['v2_in_z'],
        recovered_autoencoder_network['v2_in_dz'], 
        recovered_autoencoder_network['v2_in_sindy_predict'], 
        recovered_autoencoder_network['v2_in_decoder_weights'],
        recovered_autoencoder_network['v2_in_decoder_biases'])

# Here I access the tensors recovered from the saved Tensorflow model and convert them to torch.
converted_stack = [torch.tensor(recovered_autoencoder_network['v2_in_z'],
                                device=torch.device(params['target_device']),
                                dtype=params['sindy_precision']),
                       torch.tensor(recovered_autoencoder_network['v2_in_dz'],
                                device=torch.device(params['target_device']),
                                dtype=params['sindy_precision']),
                       torch.tensor(recovered_autoencoder_network['v2_in_sindy_predict'],
                                device=torch.device(params['target_device']),
                                dtype=params['sindy_precision'])]

# Here I use the tensors captured from the tensorflow model (converted to torch)
# with the torch version of the function and the layers from the model. 
dx_decode_torch_test, ddx_decode_torch_test = \
    calculate_first_and_second_derivative_with_torch(converted_stack,
        sindy_autoencoder.ψ_decoder_to_x) 

# Here I show the error between the two functions. 
print(dx_decode_np_test - dx_decode_torch_test, ddx_decode_np_test - ddx_decode_torch_test)

# Here I show that the Torch weights in the model feeding the Torch 
# function are equivalent to the Numpy arrays feeding the Numpy
# function. (the weights were initialized from those arrays after conversion to Torch.tensor. 
print(("\n\nWeight and bias comparison for two models (imported from np source)\n\n" +
    "weights comparison: \nl1 {:.5f} ({:.2%})\nl2 {:.5f} ({:.2%})\nl3 {:.5f} ({:.2%})\nl4 {:.5f} ({:.2%})\n\n" +
    "bias comparison: \nb1 {:.5f} ({:.2%})\nb2 {:.5f} ({:.2%})\nb3 {:.5f} ({:.2%})\nb4 {:.5f} ({:.2%}))")
        .format(np.sum(sindy_autoencoder.ψ_decoder_to_x[0].weights.cpu().detach().numpy()
                 - recovered_autoencoder_network['v2_in_decoder_weights'][0]),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[0].weights.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_weights'][0]))
                  / np.sum(recovered_autoencoder_network['v2_in_decoder_weights'][0]),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[1].weights.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_weights'][1])),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[1].weights.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_weights'][1]))
                  / np.sum(recovered_autoencoder_network['v2_in_decoder_weights'][1]),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[2].weights.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_weights'][2])),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[2].weights.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_weights'][2]))
                  / np.sum(recovered_autoencoder_network['v2_in_decoder_weights'][2]),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[3].weights.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_weights'][3])),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[3].weights.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_weights'][3]))
                  / np.sum(recovered_autoencoder_network['v2_in_decoder_weights'][3]),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[0].bias.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_biases'][0])),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[0].bias.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_biases'][0]))
                  / np.sum(recovered_autoencoder_network['v2_in_decoder_biases'][0]),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[1].bias.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_biases'][1])),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[1].bias.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_biases'][1]))
                  / np.sum(recovered_autoencoder_network['v2_in_decoder_biases'][1]),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[2].bias.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_biases'][2])),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[2].bias.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_biases'][2]))
                  / np.sum(recovered_autoencoder_network['v2_in_decoder_biases'][2]),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[3].bias.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_biases'][3])),
                  (np.sum(sindy_autoencoder.ψ_decoder_to_x[3].bias.cpu().detach().numpy()
                          - recovered_autoencoder_network['v2_in_decoder_biases'][3]))
                  / np.sum(recovered_autoencoder_network['v2_in_decoder_biases'][3])))

原因分析 #

• 硬件与底层实现差异：NumPy运行在CPU上，使用BLAS库进行矩阵运算；PyTorch默认在CUDA上运行，使用cuBLAS库。两者在浮点数舍入、运算顺序上存在细微差别，多次迭代后误差累积，导致最终结果出现显著偏差。
• Sigmoid实现不一致：PyTorch的torch.sigmoid()为了数值稳定性，在极端值场景下做了优化处理，而自定义NumPy sigmoid（1/(1+np.exp(-x))）是直接实现，两者输出存在细微差异，导数计算依赖sigmoid结果，误差被放大。
• 张量内存布局影响：CUDA张量的内存布局与NumPy的CPU数组不同，矩阵乘法的缓存优化策略可能导致计算结果的微小差异，累积后影响最终精度。

解决方案 #

1. 排除CUDA硬件影响 #

先将PyTorch运算切换到CPU，验证偏差是否由CUDA导致：

params: Dict[str, Any] = {'weight_precision': torch.float64,
                          'sindy_precision': torch.float64,
                          'target_device': 'cpu'}

2. 对齐Sigmoid实现 #

将PyTorch中的sigmoid替换为与NumPy完全一致的自定义实现：

def custom_sigmoid(x):
    return 1 / (1 + torch.exp(-x))

# 在PyTorch函数中替换：
# x = torch.sigmoid(x)
x = custom_sigmoid(x)

3. 优化导数计算步骤 #

拆分导数计算的连续运算，减少中间精度损失：

# 原PyTorch导数计算
sigmoid_first_derivative = torch.mul(x, 1 - x)
sigmoid_second_derivative = torch.mul(sigmoid_first_derivative, 1 - 2 * x)

# 修改为分步计算
one_minus_x = 1 - x
sigmoid_first_derivative = x * one_minus_x
one_minus_two_x = 1 - 2 * x
sigmoid_second_derivative = sigmoid_first_derivative * one_minus_two_x

4. 对齐矩阵乘法实现 #

使用torch.matmul替代torch.mm，与NumPy的np.matmul行为更一致：

# 原代码
x = torch.mm(x, stack[i].weights) + stack[i].bias
dx_prev = torch.mm(dx, stack[i].weights)

# 修改为
x = torch.matmul(x, stack[i].weights) + stack[i].bias
dx_prev = torch.matmul(dx, stack[i].weights)

5. 强制全局双精度 #

确保所有运算都使用float64，避免隐式类型转换：

torch.set_default_dtype(torch.float64)

内容的提问来源于stack exchange，提问作者simon.watson.sjw