File size: 4,619 Bytes
63eca6c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 |
import torch
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import seaborn as sns
import umap
import os
from post_train import collect_features, prepare_data_and_model
def visualize_features(features, labels, method='tsne', save_path=None, include_unknown=False):
"""
可视化特征分布
Args:
features: torch.Tensor, 特征向量
labels: torch.Tensor, 标签
method: str, 'tsne' 或 'umap'
save_path: str, 保存路径，如果为None则显示图像
include_unknown: bool, 是否包含未知类（第21类）
"""
# 转换为numpy数组
features = features.cpu().numpy()
labels = labels.cpu().numpy()
# 降维
print(f"Performing {method.upper()} dimensionality reduction...")
if method.lower() == 'tsne':
reducer = TSNE(n_components=2, random_state=42)
embedded = reducer.fit_transform(features)
else: # umap
reducer = umap.UMAP(n_components=2, random_state=42)
embedded = reducer.fit_transform(features)
# 清理之前的图像状态并创建新图形
plt.close('all') # 关闭所有图形
fig = plt.figure(figsize=(15, 10))
# 定义标记样式和颜色
markers = ['o', 's', '^', 'D'] # 圆形、方形、三角形、菱形
colors = ['#FF4B4B', '#4B4BFF', '#4BFF4B', '#FFB74B', '#B74BFF'] # 红、蓝、绿、橙、紫
# 确定要绘制的类别数量
num_classes = 21 if include_unknown else 20
# 为每个类别分配标记和颜色
for i in range(num_classes):
marker_idx = i % len(markers)
color_idx = i % len(colors)
mask = labels == i
if i == 20: # 未知类使用特殊标记
plt.scatter(
embedded[mask, 0],
embedded[mask, 1],
c='gray', # 使用灰色
marker='*', # 使用星形
s=150, # 稍微大一点
alpha=0.6,
label='Unknown',
edgecolors='white',
linewidth=0.5
)
else:
plt.scatter(
embedded[mask, 0],
embedded[mask, 1],
c=colors[color_idx],
marker=markers[marker_idx],
s=100,
alpha=0.6,
label=f'Class {i}',
edgecolors='white',
linewidth=0.5
)
plt.grid(True, linestyle='--', alpha=0.7)
plt.title(f'Feature Distribution ({method.upper()})', fontsize=14, pad=20)
# 调整图例
plt.legend(bbox_to_anchor=(1.05, 1),
loc='upper left',
borderaxespad=0,
ncol=1, # 使用单列显示图例
fontsize=10)
# 调整布局
plt.tight_layout()
# 先保存再显示
if save_path:
print(f"Saving plot to {save_path}")
os.makedirs(os.path.dirname(save_path), exist_ok=True)
plt.savefig(save_path, bbox_inches='tight', dpi=300, pad_inches=0.5)
plt.show()
plt.close()
if __name__ == "__main__":
# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
model, train_loader, val_loader, device = prepare_data_and_model(model_path='models/resnet50_99.92.pth', model_type='resnet50', batch_size=128)
## 加载特征
print("Collecting features of training set...")
features, labels = collect_features(
model=model,
loader=train_loader,
device=device
)
print("Collecting features of validation set...")
val_features, val_labels = collect_features(
model=model,
loader=val_loader,
device=device
)
# 可视化特征
# print("Visualizing features using t-SNE...")
# visualize_features(
# features=features,
# labels=labels,
# method='tsne',
# save_path='outputs/tsne_features.png'
# )
print("Visualizing features using UMAP...")
visualize_features(
features=features,
labels=labels,
method='umap',
save_path='outputs/resnet50_train_FeatureMap.png',
include_unknown=False
)
visualize_features(
features=val_features,
labels=val_labels,
method='umap',
save_path='outputs/resnet50_val_FeatureMap.png',
include_unknown=True
)
|