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 layerdef hook_t(module, input, output):self.output_feature = outputself.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 layerself.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 = Nonefor submodule_name, submodule in self.model.named_children():x = submodule(x)# print('-----submodule_name:',submodule_name)if submodule_name in ['avgpool']:output = xself.optimizer.zero_grad()loss = self.mse_loss(output, target)print('mse_loss:',loss)loss.backward()self.optimizer.step()