TensorFlow vs PyTorch: Memory usage

Question

I have PyTorch 1.9.0 and TensorFlow 2.6.0 in the same environment, and both recognizing the all GPUs.

I was comparing the performance of both, so I did this small simple test, multiplying large matrices (A and B, both 2000x2000) several times (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)

As a result:

PyTorch: 19.86856198310852 s
TensorFlow: 2.8338065147399902 s

I was not expecting these results, I thought the results should be similar.

Investigating the GPU performance, I noticed that both are using GPU full capacity, but PyTorch uses a small fraction of the memory Tensorflow uses. It explains the processing time difference, but I cannot explain the difference on memory usage. Is it something intrinsic to the methods, or is it my computer configuration? Regardless the matrix size (at least for matrices larger than 1000x1000), these plateau are the same.

Thanks you for your help.

user_5523920 · Answer

It is because you are doing matrix multiplication in pytorch but element-wise multiplication in tensorflow. To do matrix multiplication in TF, use tf.matmul or simply:

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

That does the same for both TF and torch. You also have to call torch.cuda.synchronize() inside the time measurement and move torch.cuda.empty_cache() outside of the measurement for the sake of fairness.

The expected results will be tensorflow's eager execution slower than pytorch.

Regarding the memory usage, TF by default claims all GPU memory and using nvidia-smi in linux or similarly task manager in windows, does not reflect the actual memory usage of the operations.

TensorFlow vs PyTorch: Memory usage

1 Answers