【t-SNE可视化CNNs特征向量-代码】

t-SNE可视化CNNs特征向量-代码

本博客主要是自己学习记录，参考网络，欢迎指正

整体代码

ModelPath是存放训练好的模型参数的路径
DatasetPath是存放数据集的文件夹的路径，其中不同类别放在不同的子文件夹里

也可以参考【t-SNE可视化-代码】

import os
import torch as t
from torch import nn
from torch.utils.data import DataLoader as DL
from torch import optim
import torchvision.datasets as datasets
from torchvision import transforms as T
import torchvision.models as models
import numpy as np
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE

from Config.config import ConfigCLS
from Classification.Dataset.ChestXRayCLS.chestxray import ChestXRay
from Classification.Dataset.RefugeCLS.refuge import REFUGE
from Classification.TrainFunc.runner import NetRunner

ModelPATH = 'resnet18/checkpoint/best_f1score.pth'
DatasetPath = 'datasets'

class Identity(nn.Module):
    def __init__(self):
        super(Identity, self).__init__()

    def forward(self, x):
        return x

def plot_embedding(data, label, title):
    """
    :param data:数据集
    :param label:样本标签
    :param title:图像标题
    :return:图像
    """
    x_min, x_max = np.min(data, 0), np.max(data, 0)
    data = (data - x_min) / (x_max - x_min)		# 对数据进行归一化处理
    fig = plt.figure()		# 创建图形实例
    ax = plt.subplot(111)		# 创建子图
    # 遍历所有样本
    for i in range(data.shape[0]):
        # 在图中为每个数据点画出标签
        if label[i] == 0:
            # colorplt = plt.cm.Set1(0.82)
            colorplt = 'brown'
            print(1)
        elif label[i] == 1:
            # colorplt = plt.cm.Set1(1.64)
            colorplt = 'red'
            print(2)
        elif label[i] == 2:
            # colorplt = plt.cm.Set1(2.46)
            colorplt = 'orangered'
            print(3)
        elif label[i] == 3:
            # colorplt = plt.cm.Set1(3.28)
            colorplt = 'orange'
            print(4)
        elif label[i] == 4:
            # colorplt = plt.cm.Set1(4.10)
            colorplt = 'chartreuse'
            print(5)
        elif label[i] == 5:
            # colorplt = plt.cm.Set1(4.92)
            colorplt = 'cyan'
            print(6)
        elif label[i] == 6:
            # colorplt = plt.cm.Set1(5.74)
            colorplt = 'lime'
            print(7)
        elif label[i] == 7:
            # colorplt = plt.cm.Set1(6.56)
            colorplt = 'dodgerblue'
            print(8)
        elif label[i] == 8:
            # colorplt = plt.cm.Set1(7.38)
            colorplt = 'blue'
            print(9)
        elif label[i] == 9:
            # colorplt = plt.cm.Set1(8.20)
            colorplt = 'm'
            print(10)
        elif label[i] == 10:
            # colorplt = plt.cm.Set1(9)
            colorplt = 'deeppink'
            print(11)
        plt.scatter(data[i, 0], data[i, 1], color=colorplt, marker='o', s=30)
    plt.xticks()		# 指定坐标的刻度
    plt.yticks()
    plt.title(title, fontsize=14)
    return fig

if __name__ == '__main__':

    img_transform_test = T.Compose([
        T.Resize((256, 256)),
        T.ToTensor(),
        T.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
    ])

    feature_list = []
    label_list = []

    test_set = datasets.ImageFolder(DatasetPath, transform=img_transform_test)
    test_loader = DL(test_set, batch_size=8, shuffle=True, num_workers=0)

    net = models.resnet18()
    net.fc = Identity()
    net = net
    net.load_state_dict(t.load(ModelPATH), False)
    net.eval()

    num = 0

    for index, data in enumerate(test_loader):
        pic, label = data

        prediction = net(pic)

        if num == 0:
            feature_list = prediction.clone().detach()
            label_list = label.clone().detach()
        else:
            feature_list = t.cat((feature_list, prediction.clone().detach()), dim=0)
            label_list = t.cat((label_list, label.clone().detach()), dim=0)

        num =+ 1

    ts = TSNE(perplexity=50, n_components=2, init='pca', random_state=0)
    result = ts.fit_transform(feature_list)
    fig = plot_embedding(result, label_list, 't-SNE Embedding of digits')

    plt.show()

结果图展示