模型参数加载中遇到的一些问题（BN层）

最新推荐文章于 2024-04-05 17:27:46 发布

Tiky0n

最新推荐文章于 2024-04-05 17:27:46 发布

阅读量836

点赞数 1

文章标签： python 开发语言

本文链接：https://blog.csdn.net/Tikyon/article/details/126141562

版权

在自定义VGG16BN模型并加载权重时遇到问题，包括BN层中未在权重文件中找到的track_running_stats参数，BN层参数名的误导，以及如何处理BN层在部署时的running_mean和running_var更新。同时讨论了PyTorch中OrderedDict在处理模型参数时的方法。

摘要由CSDN通过智能技术生成

最近需要用到VGG16BN模型，便于需求，自拟了模型，没有使用官方源码，这导致在加载参数过程中遇到一些问题：

1.自拟的模型中BN层出现了一个权重文件中没有的参数：track_running_stats:

nn.BatchNorm2d(channels, eps=1e-05, momentum=0.1, affine=False, track_running_stats=True)

这个参数能够:

①训练时统计forward过的min-batch数目,每经过一个min-batch, track_running_stats+=1；

②如果没有指定momentum, 则使用1/num_batches_tracked 作为因数来计算均值和方差(running mean and variance).

【如下所示】

class _BatchNorm(_NormBase):
    def __init__(
        self,
        num_features,
        eps=1e-5,
        momentum=0.1,  
        ##  如果不为None，会执行track_running_stats+1；如果是None，则执行通过t_r_s计算momentum
        affine=True,
        track_running_stats=True,  
        device=None,
        dtype=None
    ):
        factory_kwargs = {'device': device, 'dtype': dtype}