深入理解GPU内存分配：机器学习工程师的实用指南与实验

给定一个模型架构、数据类型、输入形状和优化器,你能否计算出前向传播和反向传播所需的GPU内存量?要回答这个问题,我们需要将流程分解为基本组件,并从底层理解内存需求。以下实验(可以在Google Colab上运行)将帮助你理解核心概念。

 

数据类型

1. float32

需要4字节的内存，

1. bfloat16

需要2字节，我们可以绘制一些数据类型所需的内存图。

图1：不同数据类型的内存分配

 

https://avoid.overfit.cn/post/2b0e1514fcb74a2c90bb5016b6b12cd9