TensorFlow 与 PyTorch：内存使用

Question

我在相同的环境中有 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 的矩阵），这些平台都是相同的。

感谢您的帮助。

Answer 1

这是因为你在pytorch中做的是矩阵乘法，而在tensorflow中做的是逐元素乘法。要在 TF 中进行矩阵乘法，请使用 tf.matmul 或简单地使用：

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

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

预期的结果是tensorflow的eager执行比pytorch慢。

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

TensorFlow 与 PyTorch：内存使用

TensorFlow vs PyTorch: Memory usage

memory-management

gpu

python-3.x

tensorflow

pytorch