怎么预估pytorch的显存占用?

参考来源:

显存占用计算

初衷:

如何让网络接收一张超大分辨图片的图片,如果直接输入会cuda of meomory,只有使用多卡模型并行的方式,就是将网络拆开放入多个GPU,但每个子网络的显存不能超过最大显存,这就需要估计各个子网络的显存占用.

环境:

ubuntu 20.04
python 3.8.19
pytorch 1.10.0 py3.8_cuda11.3_cudnn8.2.0_0

实验:

实验一:论证显存和总参数量的关系

使用nvidia-smi -l 监测显存真实占用
使用torchsummary得到模型参数总量,包括输入图片,模型参数量,中间层特征图大小,前向传播和后向传播量,好像还缺了一个优化器?

"""
使用torchinfo进行显存推断,和torchsummary区别不大
"""
from torchvision.models.resnet import ResNet, Bottleneck
from torchsummary import summary

class SingleGPUResNet50(ResNet):
    def __init__(self, *args, **kwargs):
        super(SingleGPUResNet50, self).__init__(
            Bottleneck, [3, 4, 6, 3], num_classes=1000, *<