pytorch 获取模型中间层输出的方法

最新推荐文章于 2022-11-21 17:39:17 发布

小帆97

最新推荐文章于 2022-11-21 17:39:17 发布

阅读量1.4k

点赞数 1

分类专栏： pytorch 文章标签： pytorch 深度学习 python

本文链接：https://blog.csdn.net/weixin_41396062/article/details/127617352

版权

pytorch 专栏收录该内容

1 篇文章 0 订阅

订阅专栏

1.使用hook钩子函数

获取resnet50中Avgpool层的输出，通过model.avgpool.register_forward_hook(hook_t)函数，可以获取到Avgpool层的输出

class ModelConfigure():
    def __init__(self, learning_rate = 1e-4):
        self.model = wide_resnet50_2(pretrained=pretrained_flag, progress=True)
        self.model.to('cpu')
        self.model.eval()
        
        none_layer = nn.Sequential()
        self.model.fc = none_layer  # delete fc layer
        
        def hook_t(module, input, output):
            self.output_feature = output
        self.model.avgpool.register_forward_hook(hook_t)

        self.optimizer = torch.optim.SGD(self.model.parameters(), lr=learning_rate)
        self.mse_loss = torch.nn.MSELoss()

    def model_inference_and_train(self,x,target):
        self.model.forward(x)
        self.optimizer.zero_grad()
        loss = self.mse_loss(self.output_feature, target)
        print('mse_loss:',loss)
        loss.backward()
        self.optimizer.step()

2.使用model.named_children()

通过模块的名称判断是否为想要的输出

class ModelConfigure():
    def __init__(self, learning_rate = 1e-4):

        self.model = wide_resnet50_2(pretrained=pretrained_flag, progress=True)
        self.model.to('cpu')
        self.model.eval()
        none_layer = nn.Sequential()
        self.model.fc = none_layer  # delete fc layer
        
        self.optimizer = torch.optim.SGD(self.model.parameters(), lr=learning_rate)
        self.mse_loss = torch.nn.MSELoss()
        
        
    def model_inference_and_train(self,x,target):
        output = None
        for submodule_name, submodule in self.model.named_children():
            x = submodule(x)
            # print('-----submodule_name:',submodule_name)
            if submodule_name in ['avgpool']:
                output = x
                
        self.optimizer.zero_grad()
        loss = self.mse_loss(output, target)
        print('mse_loss:',loss)
        loss.backward()
        self.optimizer.step()