TensorFlow vs PyTorch：内存使用答案

【问题标题】：TensorFlow vs PyTorch: Memory usageTensorFlow vs PyTorch：内存使用
【发布时间】：2021-10-26 07:52:54
【问题描述】：

我在同一环境中有 PyTorch 1.9.0 和 TensorFlow 2.6.0，并且都可以识别所有 GPU。

我在比较两者的性能，所以我做了这个简单的小测试，将大矩阵（A 和 B，都是 2000x2000）乘以几次（10000x）：

import numpy as np
import os
import time

def mul_torch(A,B):
    # PyTorch matrix multiplication
    os.environ['KMP_DUPLICATE_LIB_OK']='True'
    import torch
    A, B = torch.Tensor(A.copy()), torch.Tensor(B.copy())
    A = A.cuda()
    B = B.cuda()
    start = time.time()
    for i in range(10000):
        C = torch.matmul(A, B)
    torch.cuda.empty_cache()
    print('PyTorch:', time.time() - start, 's')

    return C

def mul_tf(A,B):
    # TensorFlow Matrix Multiplication
    import tensorflow as tf
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    with tf.device('GPU:0'):
        A = tf.constant(A.copy())
        B = tf.constant(B.copy())
    start = time.time()
    for i in range(10000):
        C = tf.math.multiply(A, B)
    print('TensorFlow:', time.time() - start, 's')
    return C

if __name__ == '__main__':
    A = np.load('A.npy')
    B = np.load('B.npy')

    n = 2000
    A = np.random.rand(n, n)
    B = np.random.rand(n, n)
    
    PT = mul_torch(A, B)
    time.sleep(5)
    TF = mul_tf(A, B)

结果：

PyTorch: 19.86856198310852 s
TensorFlow: 2.8338065147399902 s

我没想到这些结果，我认为结果应该是相似的。

调查GPU performance，我注意到两者都在使用 GPU 的全部容量，但 PyTorch 使用了 Tensorflow 使用的一小部分内存。它解释了处理时间差异，但我无法解释内存使用量的差异。它是方法固有的，还是我的计算机配置？无论矩阵大小（至少对于大于 1000x1000 的矩阵），这些平台都是相同的。

感谢您的帮助。

【问题讨论】：

你尝试过更大的（100000x100000）矩阵乘法吗？

标签： python-3.x tensorflow memory-management pytorch gpu

【解决方案1】：

这是因为您在 pytorch 中进行矩阵乘法，但在 tensorflow 中进行逐元素乘法。要在 TF 中进行矩阵乘法，请使用 tf.matmul 或干脆：

for i in range(10000):
  C = A @ B

这对 TF 和 torch 都是一样的。为了公平起见，您还必须在时间测量内调用torch.cuda.synchronize()，并将torch.cuda.empty_cache()移到测量外。

预期结果将是 tensorflow 的急切执行比 pytorch 慢。

关于内存使用情况，TF默认声明所有GPU内存，在linux中使用nvidia-smi或在windows中类似的任务管理器，并不反映操作的实际内存使用情况。

【讨论】：