Pytorch训练Resnet101,显存爆炸

最近需要把caffe上的resnet101网络模型迁移到pytorch上,caffe上显存还够使,一到pytorch上就不够用了,即使是resnet50也不够用。

  • 显卡:11G的1080ti
  • batch_size:1
  • image_size:480
    两天各种调试,最终还是在网络层代码中发现了问题:

    很明显的planes=2056应该是planes=256!
    这个错误很低级。
模型大小
  • 修改之前,pytorch上resnet50模型大小
  • 修改之前,pytorch上resnet101模型大小
  • 修改之后,pytorch上resnet50模型大小
  • 修改之后,pytorch上resnet101模型大小
  • 迁移前,caffe上resnet101模型大小
显存消耗
  • 修改之前,pytorch上resnet50显存消耗
    超出显存大小
  • 修改之前,pytorch上resnet101显存消耗
    超出显存大小
  • 修改之后,pytorch上resnet50显存消耗
  • 修改之后,pytorch上resnet101显存消耗

ModuleList和Sequential区别

网上很多人写了二者区别,我就说下自己的感受吧,先看代码:

def __init__(self):
        layers5 = []
        for i in range(0, 3):
                layers5.append(Bottleneck2(inplanes=2048, planes=512))
        self.layer5 = nn.Sequential(*layers5)
def forward(self, x):
        x = self.layer5(x)
def __init__(self):
        self.fiveth_stage = nn.ModuleList()
        for i in range(0, 3):
                self.fiveth_stage.append(Bottleneck2(inplanes=2048, planes=512))
def forward(self, x):
        for layer in self.fiveth_stage:
                x = layer(x)

很明显Sequential会对其里面的每个操作自动forward,而ModuleLis只是封装的list,需要遍历list前向传播.

你可能感兴趣的:(Pytorch训练Resnet101,显存爆炸)