自动微分对比：PyTorch与TensorFlow在GBM模拟中的内存性能差异

我实现了一个模拟N条几何布朗运动（GBM）路径的函数，包含M个离散点（含起点共M+1个），函数返回GBM最后一步的取值（完整路径会保留以备后续使用）。我需要通过自动微分（autograd、算法微分AAD等）获取输出向量对输入spot和vol的导数，即尺寸为N×2的雅可比矩阵J。

我分别基于PyTorch和TensorFlow实现了逻辑相近的函数以对比性能，代码如下：

PyTorch实现 #

import torch
from torch.autograd.functional import jacobian


def sim_gbm(N, t, spot, drift, vol, seed=None):
    M = int(t.shape[0])
    dt = torch.diff(t)

    Z = torch.normal(mean=0.0, std=1.0, size=(M - 1, N), generator=torch.manual_seed(seed))

    W = torch.concatenate([
        torch.zeros(size=(1, N)), torch.sqrt(dt)[:, None] * Z
    ]).cumsum(axis=0)

    S = spot * torch.exp(((drift - 0.5 * vol ** 2) * t)[:, None] + vol * W)
    return S[-1,]


def gbm_wrapper(spot, vol):
    return sim_gbm(N, t, spot, drift, vol, seed)


if __name__ == '__main__':
    from datetime import datetime
    seed = 1234
    N = 10000
    M = 52
    t0 = 0.0
    T = 1.0
    spot = torch.tensor(100.0, requires_grad=True)
    drift = torch.tensor(0.03)
    vol = torch.tensor(0.2, requires_grad=True)
    t = torch.linspace(t0, T, M + 1)

    start = datetime.now()

    J = jacobian(func=gbm_wrapper, inputs=(spot, vol))
    stop = datetime.now()
    print(stop - start)

TensorFlow实现 #

import tensorflow as tf

def sim_gbm(N, t, spot, drift, vol, seed=None):
    M = tf.constant(t.shape[0])
    dt = tf.math.subtract(t[1:], t[:-1])

    Z = tf.random.normal(mean=0.0, stddev=1.0, shape=(M - 1, N), seed=seed)

    W = tf.concat([
        tf.zeros(shape=(1, N)), tf.sqrt(dt)[:, None] * Z
    ], axis=0)

    S = spot * tf.exp(((drift - 0.5 * vol ** 2) * t)[:, None] + vol * W)
    return S[-1,]


if __name__ == '__main__':
    from datetime import datetime
    seed = tf.constant(1234)
    N = tf.constant(10000)
    M = 52
    t0 = tf.constant(0.0)
    T = tf.constant(1.0)
    spot = tf.constant(100.0, dtype=tf.float32)
    drift = tf.constant(0.03, dtype=tf.float32)
    vol = tf.constant(0.2, dtype=tf.float32)
    t = tf.linspace(t0, T, M + 1)

    start = datetime.now()
    with tf.GradientTape() as tape:
        tape.watch([spot, vol])
        S = sim_gbm(N, t, spot, drift, vol, seed)

    J = tape.jacobian(S, [spot, vol])

    stop = datetime.now()
    print(stop - start)

测试结果 #

• 当N≤1000时，两者性能相近，耗时均约1秒；但PyTorch仅占用约140Mb内存，TensorFlow则占用约340Mb内存。
• 当N增大至10000时，性能差异显著：PyTorch耗时约1分钟，内存仍维持在140Mb；而TensorFlow运行数分钟后因耗尽16Gb内存报错。

疑问 #

既然两个框架在处理含数百万参数的神经网络时都能应对更大的计算图，为何在此场景下会出现如此大的内存差异？

内容的提问来源于stack exchange，提问作者Landscape