预训练时候 Pytorch FrozenBatchNorm (BN)一些小心得

最新推荐文章于 2024-04-20 19:38:36 发布

置顶 zqx951102

最新推荐文章于 2024-04-20 19:38:36 发布

阅读量1.9k

点赞数 3

分类专栏：算法 Python 论文文章标签： pytorch 深度学习 python

本文链接：https://blog.csdn.net/zqx951102/article/details/128346879

版权

Python 同时被 3 个专栏收录

70 篇文章 13 订阅

订阅专栏

论文

43 篇文章 6 订阅

订阅专栏

算法

32 篇文章 4 订阅

订阅专栏

在预训练模型中会发现这样使用：

# resnet model builder function
def build_resnet(arch='resnet50', pretrained=True,
        freeze_backbone_batchnorm=True, freeze_layer1=True,
        norm_layer=misc_nn_ops.FrozenBatchNorm2d):
    # weights
    if pretrained: #如果是预训练 权重是xxx
        weights = torchvision.models.ResNet50_Weights.IMAGENET1K_V1
    else:
        weights = None

    # load model
    if freeze_backbone_batchnorm:  #如果需要冻结BN层 那么就选择 normlayer为FrozenBatchNorm2d
        resnet = torchvision.models.resnet50(weights=weights, norm_layer=norm_layer)
    else:
        resnet = torchvision.models.resnet50(weights=weights)

    # freeze first layers
    resnet.conv1.requires_grad_(False)
    resnet.bn1.requires_grad_(False)
    if freeze_layer1:
        resnet.layer1.requires_grad_(False)

    # setup backbone architecture
    backbone, head = ResnetBackbone(resnet), ResnetHead(resnet)
    
    # return backbone, head
    return backbone, head

也就是需要设置bn为norm_layer=misc_nn_ops.FrozenBatchNorm2d，冻结住bn的参数，这样做的好处有：

FrozenBatchNorm就是"weight" and “bias”, “running_mean”, "running_var”四个值固定住的BN

经典框架中一直使用的是FrozenBatchNorm2d。如Detectron，DETR, mmdetection？见

“weight” and “bias”, “running_mean”, "running_var”四个值是buf，通过register_buffer设置不更新。

为什么要使用FrozenBatchNorm
BN层在CNN网络中大量使用，但是BN依赖于均值和方差，如果batch_size太小，计算一个小batch_size的均值和方差，肯定没有计算大的batch_size的均值和方差稳定和有意义，这个时候，还不如不使用bn层，因此可以将bn层冻结。另外，我们使用的网络，几乎都是在imagenet上pre-trained，完全可以使用在imagenet上学习到的参数。

而且，如果使用的是FrozenBatchNorm，多卡训练就不会有BN同步的问题了，那么多卡训练的性能理论上应该和单卡一样好了，注意这点

torchvision.ops.FrozenBatchNorm2d(num_features: int, eps: float = 1e-05)
在这里插入图片描述