pytorch 对加载好的模型重新进行初始化

最新推荐文章于 2024-06-21 14:35:06 发布

原创

最新推荐文章于 2024-06-21 14:35:06 发布 · 1.8k 阅读

6 ·

CC 4.0 BY-SA版权

这是一个用于PyTorch模型权重初始化的脚本，涵盖了多种层类型，如Conv1d、Conv2d、Conv3d、Linear、LSTM等，并使用了不同的初始化方法，如normal_、xavier_normal_、orthogonal_。通过调用apply(weight_init)方法，可以方便地为整个模型的权重和偏置进行重置。

在验证小模型从头训练的效果时，需要把已有的权重重新进行初始化，发现了一个封装好的功能脚本，感谢大佬。

#!/usr/bin/env python
# -*- coding:UTF-8 -*-

import torch
import torch.nn as nn
import torch.nn.init as init


def weight_init(m):
    '''
    Usage:
        model = Model()
        model.apply(weight_init)
    '''
    if isinstance(m, nn.Conv1d):
        init.normal_(m.weight.data)
        if m.bias is not None:
            init.normal_(m.bias.data)
    elif isinstance(m, nn.Conv2d):
        init.xavier_normal_(m.weight.data)
        if m.bias is not None:
            init.normal_(m.bias.data)
    elif isinstance(m, nn.Conv3d):
        init.xavier_normal_(m.weight.data)
        if m.bias is not None:
            init.normal_(m.bias.data)
    elif isinstance(m, nn.ConvTranspose1d):
        init.normal_(m.weight.data)
        if m.bias is not None:
            init.normal_(m.bias.data)
    elif isinstance(m, nn.ConvTranspose2d):
        init.xavier_normal_(m.weight.data)
        if m.bias is n