【发布时间】：2022-01-01 23:42:43
【问题描述】：
我已经使用 TensorFlow 测试了 CPU 到 GPU 的数据传输吞吐量，它似乎明显低于 PyTorch。对于慢 2 倍和 5 倍之间的大张量。在 TF 中，我达到了 25MB 张量（~4 GB/s）的最大速度，随着张量大小的增加，它下降到 2 GB/s。 PyTorch 数据传输速度随着张量大小而增长，并在 9 GB/s（25MB 张量）时达到饱和。该行为在 RTX 2080ti 和 GTX 1080ti 以及 TF 2.4 和 2.6 上是一致的。
我做错了吗？有什么方法可以匹配 PyTorch 的数据吞吐量？我不只是想隐藏延迟，例如使用异步队列，但我想获得完整的数据带宽。
TF 中批量 256x256x3 图像的结果（平均超过 100 次传输）：
code: tf.cast(x, dtype=tf.float32)[0, 0]
Batch size 1; Batch time 0.0005; BPS 1851.8; FPS 1851.8; MB/S 364.1
Batch size 2; Batch time 0.0004; BPS 2223.5; FPS 4447.1; MB/S 874.3
Batch size 4; Batch time 0.0006; BPS 1555.2; FPS 6220.6; MB/S 1223.0
Batch size 8; Batch time 0.0006; BPS 1784.8; FPS 14278.7; MB/S 2807.3
Batch size 16; Batch time 0.0013; BPS 755.3; FPS 12084.7; MB/S 2376.0
Batch size 32; Batch time 0.0023; BPS 443.8; FPS 14201.3; MB/S 2792.1
Batch size 64; Batch time 0.0035; BPS 282.5; FPS 18079.5; MB/S 3554.6
Batch size 128; Batch time 0.0061; BPS 163.4; FPS 20916.4; MB/S 4112.3
Batch size 256; Batch time 0.0241; BPS 41.5; FPS 10623.0; MB/S 2088.6
Batch size 512; Batch time 0.0460; BPS 21.7; FPS 11135.8; MB/S 2189.4
与 PyTorch 的结果相同：
Code: torch.from_numpy(x).to(self.device).type(torch.float32)[0, 0].cpu()
Batch size 1; Batch time 0.0001; BPS 10756.6; FPS 10756.6; MB/S 2114.8
Batch size 1; Batch time 0.0001; BPS 12914.7; FPS 12914.7; MB/S 2539.1
Batch size 2; Batch time 0.0001; BPS 10204.4; FPS 20408.7; MB/S 4012.5
Batch size 4; Batch time 0.0002; BPS 5841.1; FPS 23364.3; MB/S 4593.6
Batch size 8; Batch time 0.0003; BPS 3994.4; FPS 31955.4; MB/S 6282.7
Batch size 16; Batch time 0.0004; BPS 2713.8; FPS 43421.3; MB/S 8537.0
Batch size 32; Batch time 0.0007; BPS 1486.3; FPS 47562.7; MB/S 9351.2
Batch size 64; Batch time 0.0015; BPS 679.3; FPS 43475.9; MB/S 8547.7
Batch size 128; Batch time 0.0028; BPS 359.5; FPS 46017.7; MB/S 9047.5
Batch size 256; Batch time 0.0054; BPS 185.2; FPS 47404.1; MB/S 9320.0
Batch size 512; Batch time 0.0108; BPS 92.9; FPS 47564.5; MB/S 9351.6
重现测量的完整代码是：
import time
import numpy as np
import tensorflow as tf
import torch
import argparse
def parseargs():
parser = argparse.ArgumentParser(usage='Test GPU transfer speed in TensorFlow(default) and Pytorch.')
parser.add_argument('--pytorch', action='store_true', help='Use PyTorch instead of TensorFlow')
args = parser.parse_args()
return args
class TimingModelTF(tf.keras.Model):
def __init__(self, ):
super(TimingModelTF, self).__init__()
@tf.function
def call(self, x):
return tf.cast(x, dtype=tf.float32)[0, 0]
class TimingModelTorch(torch.nn.Module):
def __init__(self, ):
super(TimingModelTorch, self).__init__()
self.device = torch.device('cuda')
def forward(self, x):
with torch.no_grad():
return torch.from_numpy(x).to(self.device).type(torch.float32)[0, 0].cpu()
if __name__ == '__main__':
args = parseargs()
width = 256
height = 256
channels = 3
iterations = 100
model = TimingModelTorch() if args.pytorch else TimingModelTF()
for batch_size in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]:
img = np.random.randint(5, size=(batch_size, height, width, channels), dtype=np.uint8)
result = model(img)
result.numpy()
start = time.time()
for i in range(iterations):
result = model(img)
result.numpy()
batch_time = (time.time() - start) / iterations
print(f'Batch size {batch_size}; Batch time {batch_time:.4f}; BPS {1 / batch_time:.1f}; FPS {(1 / batch_time) * batch_size:.1f}; MB/S {(((1 / batch_time) * batch_size) * 256 * 256 * 3) / 1000000:.1f}')
【问题讨论】：
-
可能 pytorch 使用固定缓冲区和 tensorflow 仍然可以流水线化多个操作以接近固定缓冲区的性能。
-
我不确定我是否理解。该代码不使用固定内存（主机） - 它是一个绝对分页的 numpy 数组。流水线将如何提高 CPU-GPU 吞吐量？我对固定内存的理解来自developer.nvidia.com/blog/how-optimize-data-transfers-cuda-cc
-
将数组固定到 gpu，而不是 cpu，应该减少 tf.对于 pytorch，如果 .cpu() 已经在 cpu 中，则 .cpu() 会返回没有副本的原始对象。
-
好的。固定到 GPU = 将所有数据复制到 GPU 并保留在那里并仅使用该数据。这本身无济于事，数据不适合 GPU 内存。问题仍然存在 - 我可以比发布的代码更快地将数据传输到 GPU 吗？在代码中 .cpu() 用于将数据从设备返回到主机 - 我不明白相关评论。
标签： python performance tensorflow gpu pci-e