TensorFlow与PyTorch显存占用差异及矩阵运算性能对比疑问

核心逻辑错误 #

你编写的测试代码调用了两类完全不等价的运算，这是性能、显存差异的根本原因：

• PyTorch侧调用的torch.matmul(A, B)为矩阵乘法（点积），对尺寸为n×n的输入，单次运算时间复杂度为O(n³)，运算过程中需要额外的临时存储空间存储分块运算的中间结果，计算量和显存开销都远高于逐元素运算。
• TensorFlow侧调用的tf.math.multiply(A, B)为逐元素乘法（哈达玛积），单次运算时间复杂度为O(n²)，仅需要和输入矩阵同尺寸的空间存储输出结果，计算量比矩阵乘法小两个数量级。

将TensorFlow侧的运算替换为tf.matmul(A, B)后重新测试，两者的运算速度、实际显存占用就会基本对齐。

显存平台效应说明 #

你观察到的大尺寸矩阵下显存占用达到固定平台值，是两个框架默认的显存分配策略差异导致的，和硬件配置无关：

• TensorFlow 2.x默认会预占当前设备几乎全部可用GPU显存供进程独享，避免后续动态分配显存的开销，所以你看到的显存占用是预占的总量，而非运算实际使用的量。
• PyTorch默认采用按需分配的显存策略，仅会申请当前运算实际需要的显存空间，所以你看到的占用数值就是运算本身的真实显存开销。你也可以通过调整PYTORCH_CUDA_ALLOC_CONF环境变量修改PyTorch的显存分配规则，实现类似TensorFlow的预占效果。

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