torch获取每层权重

pytorch读取每层权重

import torch
from models import Net
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

pth_path='epoch-best.pth' 
model_demo = Net()  # 构造模型，创建新模型，网络实例
loaded_model = torch.load(pth_path)  # 加载模型参数
model_demo.load_state_dict(loaded_model['state_dict'])  #将模型参数加载到构造的新建模型实例model_demo中，需要创建的model_demo模型和加载模型的结构、参数名称、参数维度相同，不同时，可选择加载相同部分参数
model_demo = model_demo.to(device)
model_demo.eval()

dummy_input = torch.autograd.Variable(torch.randn((1,3,320,320)))
torch_outputs=model(dummy_input)
# style1
for name in model.state_dict():                                             
    print('model weight name ',name)                                        
    if name=='head.gfl_cls.2.weight':                                       
        print('head.gfl_cls.2.weight:',model.state_dict()[name])
# style2             
for name, parameters  in model.state_dict.items(): 
    print('model weight name ',name)
    if name=='head.gfl_cls.2.weight':  
        print(name,':',parameters.detach().numpy())
# style3        
for name, parameter in model.named_parameters():  
    if 'head.gfl_cls.2.weight' in name:
        print(name,':',parameter,parameter.shape,parameter.size())