import torch
import torch.nn as nn
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
from torchvision import transforms# 1. 数据预处理：加载并转换图像
def load_image(image_path):"""将图像转为PyTorch张量，添加批次维度"""# 预处理流程：调整尺寸→转Tensor→标准化（便于网络训练）transform = transforms.Compose([transforms.Resize((224, 224)),  # 统一尺寸为224×224transforms.ToTensor(),          # 转为Tensor，形状[C, H, W]，值范围[0,1]transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 标准化到[-1,1]])image = Image.open(image_path).convert('RGB')  # 打开图像并转为RGB三通道image_tensor = transform(image).unsqueeze(0)   # 添加批次维度，形状变为[1, 3, 224, 224]return image, image_tensor# 2. 定义卷积网络（含2个卷积层）
class ConvDemo(nn.Module):def __init__(self):super(ConvDemo, self).__init__()# 第1个卷积层：保持尺寸不变self.conv1 = nn.Conv2d(in_channels=3,         # 输入：RGB三通道out_channels=16,       # 输出：16个特征图（16个卷积核）kernel_size=3,         # 3×3卷积核stride=1,              # 步长1padding=1              # padding=1：保持输出尺寸=输入尺寸)# 第2个卷积层：尺寸减半self.conv2 = nn.Conv2d(in_channels=16,        # 输入：上一层的16个通道out_channels=8,        # 输出：8个特征图kernel_size=5,         # 5×5卷积核stride=2,              # 步长2：输出尺寸减半padding=2              # padding=2：配合步长2保持尺寸比例)def forward(self, x):x1 = self.conv1(x)  # 经过第1层卷积的输出x2 = self.conv2(x1) # 经过第2层卷积的输出return x1, x2# 3. 可视化函数：展示卷积核与特征图
def visualize_results(original_image, conv1_output, conv2_output, conv1_kernels):plt.figure(figsize=(16, 12))  # 设置画布大小# 子图1：原始图像plt.subplot(3, 1, 1)plt.title('Original Image')plt.imshow(original_image)plt.axis('off')  # 关闭坐标轴# 子图2：第1层卷积核（前4个）plt.subplot(3, 2, 3)plt.title('First 4 Kernels of Conv1 (3×3)')kernels = conv1_kernels.detach().numpy()  # 转为numpy数组（脱离计算图）kernel_grid = np.zeros((3*3, 3*4))  # 构建3×3核的展示网格（4个核）for i in range(4):kernel = kernels[i].transpose(1, 2, 0)  # 转为[H, W, C]格式（便于显示）kernel = (kernel - kernel.min()) / (kernel.max() - kernel.min())  # 归一化到[0,1]kernel_grid[:, i*3:(i+1)*3] = kernel.reshape(9, 3)  # 排列到网格plt.imshow(kernel_grid, cmap='gray')plt.axis('off')# 子图3：第1层特征图（前4个）plt.subplot(3, 2, 4)plt.title('First 4 Feature Maps of Conv1')feat1 = conv1_output[0].detach().numpy()  # 取第1个样本的特征图feat1_grid = np.zeros((224, 224*4))  # 构建特征图展示网格for i in range(4):feat = feat1[i]feat = (feat - feat.min()) / (feat.max() - feat.min())  # 归一化feat1_grid[:, i*224:(i+1)*224] = feat  # 排列到网格plt.imshow(feat1_grid, cmap='gray')plt.axis('off')# 子图4：第2层特征图（前4个）plt.subplot(3, 1, 3)plt.title('First 4 Feature Maps of Conv2 (Downsampled)')feat2 = conv2_output[0].detach().numpy()feat2_grid = np.zeros((112, 112*4))  # 步长2，尺寸变为112×112for i in range(4):feat = feat2[i]feat = (feat - feat.min()) / (feat.max() - feat.min())  # 归一化feat2_grid[:, i*112:(i+1)*112] = feat  # 排列到网格plt.imshow(feat2_grid, cmap='gray')plt.axis('off')plt.tight_layout()  # 自动调整子图间距plt.show()# 4. 主函数：执行流程
if __name__ == '__main__':# 加载图像（替换为你的图像路径，如'cat.jpg'）image, image_tensor = load_image('example.jpg')print(f"输入图像形状: {image_tensor.shape}")  # 输出：[1, 3, 224, 224]# 初始化模型并执行前向传播model = ConvDemo()conv1_out, conv2_out = model(image_tensor)  # 得到两层卷积的输出# 打印输出形状（验证尺寸计算）print(f"Conv1输出形状: {conv1_out.shape}")  # 输出：[1, 16, 224, 224]（尺寸不变）print(f"Conv2输出形状: {conv2_out.shape}")  # 输出：[1, 8, 112, 112]（尺寸减半）# 获取第1层卷积核（权重参数）conv1_kernels = model.conv1.weight  # 形状：[16, 3, 3, 3]（16个3×3×3的核）# 可视化结果visualize_results(original_image=image,conv1_output=conv1_out,conv2_output=conv2_out,conv1_kernels=conv1_kernels)