pytorch 输出中间层特征的实例

import torchimport torchvisionimport numpy as npfrom torch import nnfrom torch.nn import initfrom torch.autograd import Variablefrom torch.utils import dataEPOCH=20BATCH_SIZE=64LR=1e-2train_data=torchvision.datasets.MNIST(root='./mnist',train=True,                   transform=torchvision.transforms.ToTensor(),download=False)train_loader=data.DataLoader(train_data,batch_size=BATCH_SIZE,shuffle=True)test_data=torchvision.datasets.MNIST(root='./mnist',train=False)test_x=Variable(torch.unsqueeze(test_data.test_data,dim=1).type(torch.FloatTensor)).cuda()/255test_y=test_data.test_labels.cuda()class CNN(nn.Module):  def __init__(self):    super().__init__()    self.conv1=nn.Sequential(        nn.Conv2d(in_channels=1,out_channels=16,kernel_size=4,stride=1,padding=2),        nn.ReLU(),        nn.MaxPool2d(kernel_size=2,stride=2))    self.conv2=nn.Sequential(nn.Conv2d(16,32,4,1,2),nn.ReLU(),nn.MaxPool2d(2,2))    self.out=nn.Linear(32*7*7,10)      def forward(self,x):    per_out=[] ############修改处##############    x=self.conv1(x)    per_out.append(x) # conv1    x=self.conv2(x)    per_out.append(x) # conv2    x=x.view(x.size(0),-1)    output=self.out(x)    return output,per_out  cnn=CNN().cuda() # or cnn.cuda()optimizer=torch.optim.Adam(cnn.parameters(),lr=LR)loss_func=nn.CrossEntropyLoss().cuda()############################for epoch in range(EPOCH):  for step,(x,y) in enumerate(train_loader):    b_x=Variable(x).cuda()# if channel==1 auto add c=1    b_y=Variable(y).cuda()#    print(b_x.data.shape)    optimizer.zero_grad()    output=cnn(b_x)[0] ##原先只需要cnn(b_x) 但是现在需要用到第一个返回值##    loss=loss_func(output,b_y)# Variable need to get .data    loss.backward()    optimizer.step()        if step%50==0:      test_output=cnn(test_x)[0]      pred_y=torch.max(test_output,1)[1].cuda().data.squeeze()      '''      why data ,because Variable .data to Tensor;and cuda() not to numpy() ,must to cpu and to numpy       and .float compute decimal      '''      accuracy=torch.sum(pred_y==test_y).data.float()/test_y.size(0)      print('EPOCH: ',epoch,'| train_loss:%.4f'%loss.data[0],'| test accuracy:%.2f'%accuracy)    #                       loss.data.cpu().numpy().item() get one value  torch.save(cnn.state_dict(),'./model/model.pth')##输出中间层特征，根据索引调用##conv1: conv1=cnn(b_x)[1][0]conv2: conv2=cnn(b_x)[1][1]

res=torchvision.models.resnet18()def get_features_hook(self, input, output):# self 代表类模块本身  print(output.data.cpu().numpy().shape)handle=res.layer2.register_forward_hook(get_features_hook)a=torch.ones([1,3,224,224])