pytorch visualizes the features of each layer of CNN

In PyTorch, you can use torchvision.utils.make_grid to visualize the feature map as a grid. The specific steps are as follows:
1. Define a dataset and load the data

import torchvision.datasets as datasets
import torchvision.transforms as transforms

# 定义数据集
train_dataset = datasets.CIFAR10(root='./data', train=True, 
                                 transform=transforms.ToTensor(), download=True)
# 加载数据
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)

2. Load the pre-trained model and define a function to extract the features of each layer

import torch.nn as nn
import torchvision.models as models

# 加载预训练模型
model = models.resnet18(pretrained=True)

# 定义一个函数，用于提取每一层的特征
def get_features(x, model, layers):
    features = []
    for name, module in model._modules.items():
        x = module(x)
        if name in layers:
            features.append(x)
    return features

3. Obtain the name of a specific layer and input a batch of data from the dataset into the model to obtain the features of the corresponding layer

# 获取模型中所有层的名称
all_layers = []
for name, layer in model.named_modules():
    all_layers.append(name)

# 获取需要可视化的层的名称
layers = all_layers[4:9]

# 获取一批数据
data, _ = next(iter(train_loader))

# 将数据输入到模型中，并获取对应层的特征
features = get_features(data, model, layers)

4. Visualize features as meshes

import matplotlib.pyplot as plt
import numpy as np

# 定义一个函数，用于将特征可视化为网格
def visualize_features(features):
    nrow = len(features)
    ncols = features[0].shape[1]
    fig, axs = plt.subplots(nrow, ncols, figsize=(10, 10))

    for i in range(nrow):
        for j in range(ncols):
            img = features[i][0][j].detach().numpy()
            img = np.transpose(img, (1, 2, 0))
            img = (img - img.min()) / (img.max() - img.min())
            axs[i][j].imshow(img)
            axs[i][j].axis('off')
            if j == 0:
                axs[i][j].set_title(layers[i])
    plt.show()

# 可视化特征
visualize_features(features)

The above code will visualize the extracted features of a specific layer as a grid, and display the name of the corresponding layer on the left side of the grid.

pytorch visualizes the features of each layer of CNN

おすすめ