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()