pytorch 获取层权重，对特定层注入hook，提取中间层输出

左手的ㄟ右手 2021-06-24 15:59 490阅读 0赞

#获取模型权重
    for k, v in model_2.state_dict().iteritems():
        print("Layer {}".format(k))
        print(v)

#获取模型权重
    for layer in model_2.modules():
       if isinstance(layer, nn.Linear):
            print(layer.weight)

#将一个模型权重载入另一个模型
    model = VGG(make_layers(cfg['E']), **kwargs)
    if pretrained:
        load = torch.load('/home/huangqk/.torch/models/vgg19-dcbb9e9d.pth')
        load_state = {k: v for k, v in load.items() if k not in ['classifier.0.weight', 'classifier.0.bias', 'classifier.3.weight', 'classifier.3.bias', 'classifier.6.weight', 'classifier.6.bias']}
        model_state = model.state_dict()
        model_state.update(load_state)
        model.load_state_dict(model_state)
    return model

# 对特定层注入hook
    def hook_layers(model):
        def hook_function(module, inputs, outputs):
            recreate_image(inputs[0])
    
        print(model.features._modules)
        first_layer = list(model.features._modules.items())[0][1]
        first_layer.register_forward_hook(hook_function)

#获取层
    x = someinput
    for l in vgg.features.modules():
      x = l(x)
    modulelist = list(vgg.features.modules())
    for l in modulelist[:5]:
      x = l(x)
    keep = x
    for l in modulelist[5:]:
      x = l(x)

# 提取vgg模型的中间层输出
    # coding:utf8
    import torch
    import torch.nn as nn
    from torchvision.models import vgg16
    from collections import namedtuple
    
    
    class Vgg16(torch.nn.Module):
        def __init__(self):
            super(Vgg16, self).__init__()
            features = list(vgg16(pretrained=True).features)[:23]
            # features的第3，8，15，22层分别是: relu1_2,relu2_2,relu3_3,relu4_3
            self.features = nn.ModuleList(features).eval()
    
        def forward(self, x):
            results = []
            for ii, model in enumerate(self.features):
                x = model(x)
                if ii in {
       3, 8, 15, 22}:
                    results.append(x)
    
            vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3'])
            return vgg_outputs(*results)