使用pytorch保存模型参数与加载模型参数的方法可以对某一层的参数进行替换、复制:
打印当前网络的参数:nn.Model.state_dict()
加载网络参数:nn.Model.load_state_dict()
代码示例
mport torch.nn as nn
class Net1(nn.Module):
def __init__(self, input_dim, output_dim):
super().__init__()
self.linear1 = nn.Linear(input_dim, output_dim) # 输入的个数,输出的个数
self.linear2 = nn.Linear(output_dim, 2)
def forward(self, x):
out = self.linear1(x)
out = self.linear2(out)
return out
class Net2(nn.Module):
def __init__(self, input_dim, output_dim):
super().__init__()
self.linear = nn.Linear(input_dim, output_dim) # 输入的个数,输出的个数
def forward(self, x):
out = self.linear(x)
return out
if __name__ == '__main__':
net1 = Net1(20, 2) # 模型1
net2 = Net2(20, 2) # 模型2
# 这里可以打断点,debug一下确认参数是否变更哟
net2.linear.load_state_dict(net1.linear1.state_dict()) # 把模型2的fc层参数替换为模型1的linear1参数
更多《计算机视觉与图形学》知识,可关注下方公众号:
参考文章
pytorch获取模型某一层参数名及参数值方式:https://www.cnblogs.com/expttt/p/12461775.html
In Pytorch, what is the most efficient way to copy the learned params of a model as the initialization for a second model of the same architecture?:https://stackoverflow.com/questions/59183865/in-pytorch-what-is-the-most-efficient-way-to-copy-the-learned-params-of-a-model