PyTorch提取中間層特征

如何得到中間層特征：

如果只想得到中間層特征，而不需要得到gradient之類的，那么不需要hook函數這么復雜。只需要在forward函數中添加一行代碼，將feature賦值給self變量即可，即self.feature_map = feature

給一個例子：

# Define a Convolutional Neural Network
class 
Net(nn.Module):
    
    def __init__(self, kernel_size=5, n_filters=16, n_layers=3):
        xxx
    def forward(self, x): 
        x = self.body(self.head(x))

        self.featuremap1 = x.detach() # 核心代碼

        return F.relu(self.fc(x))

model_ft = Net()
train_model(model_ft)
feature_output1 = model_ft.featuremap1.transpose(1,0).cpu()

這樣就得到了feature_map，並保存到了feature_output變量中。

 

如何顯示中間層特征：

給出一個簡單顯示代碼

def feature_imshow(inp, title=None):
    
    """Imshow for Tensor."""
    
    inp = inp.detach().numpy().transpose((1, 2, 0))
    
    mean = np.array([0.5, 0.5, 0.5])
    
    std = np.array([0.5, 0.5, 0.5])
    
    inp = std * inp + mean
    
    inp = np.clip(inp, 0, 1)
    
    plt.imshow(inp)
    
    if title is not None:
        plt.title(title)
    
    plt.pause(0.001)  # pause a bit so that plots are updated


out = torchvision.utils.make_grid(feature_ouput1)
feature_imshow(out)

結果圖如下：