| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061 |
- import torch
- from net import FullyConnectedNet
- from PIL import Image
- import numpy as np
- from torchvision import transforms
- import torch.nn.functional as F
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
- model = FullyConnectedNet()
- model_path = "D:/code/xwdev/mlp_cro/model/mnist_net_best.pth"
- # 加载训练好的权重到模型中：
- # 1. torch.load(model_path, map_location=device) 从磁盘读取权重字典（state_dict），
- # map_location=device 保证权重被映射到当前使用的设备上（防止在 GPU 上训练的模型在纯 CPU 机器上加载失败）
- # 2. model.load_state_dict(...) 把读取到的权重参数逐个复制进模型的对应层中
- model.load_state_dict(torch.load(model_path, map_location=device))
- model.to(device) # 把模型的所有参数和缓冲区移动到指定设备上（GPU 或 CPU），之后前向计算就在该设备上进行
- model.eval()
- # 读取图片并进行预处理
- image_path = "D:/code/xwdev/mlp_cro/mnist_pytorch_samples/image_0000_label_7.png"
- image = Image.open(image_path) # 用 Pillow 打开图片文件，得到一个 PIL Image 对象
- # 定义图像预处理流水线，transforms.Compose 会把多个变换按顺序串联起来依次执行
- transform = transforms.Compose([
- transforms.Grayscale(num_output_channels=1), # 转为灰度图：无论原图是 RGB 还是 RGBA，都强制变成单通道灰度图（MNIST 是灰度图，通道数必须为 1）
- transforms.Resize((28, 28)), # 调整大小为 28x28：MNIST 模型训练时输入就是 28x28，推理时尺寸必须一致
- transforms.ToTensor(), # 转为 tensor：把 PIL 图片转成 PyTorch 张量，形状为 [C, H, W]（这里是 [1, 28, 28]），
- # 同时把像素值从 0~255 的整数缩放到 0.0~1.0 的浮点数
- ])
- input_tensor = transform(image) # 对图片执行上面定义的预处理流水线，得到形状为 [1, 28, 28] 的张量（通道、高、宽）
- input_tensor = input_tensor.unsqueeze(0) # 添加 batch 维度：在第 0 维插入一个大小为 1 的维度，
- # 形状从 [1, 28, 28] 变为 [1, 1, 28, 28]，即 NCHW 格式（批量数、通道、高、宽），
- # 因为模型要求输入是一“批”图片，即使只有一张也要包成 batch
- input_tensor = input_tensor.to(device) # 把输入张量移动到与模型相同的设备上，设备和模型必须一致，否则计算时会报错
- # 使用模型进行预测
- with torch.no_grad(): # 上下文管理器：在其内部所有操作都不会记录梯度信息。
- # 推理阶段不需要反向传播，关闭梯度可以节省内存、加快计算
- output = model(input_tensor) # 前向传播：把图片张量送入模型，得到输出。
- output = F.softmax(output, dim=1) # Softmax(dim=1)，output 是形状 [1, 10] 的概率分布，
- # 每一维表示该图片属于数字 0~9 的概率，10 个概率之和为 1
-
- predicted_class = torch.argmax(output, dim=1).item() # 取概率最大的那一维的索引作为预测类别：
- # argmax(output, dim=1) 在第 1 维（类别维）上找最大值的下标，返回形状 [1] 的张量
- # .item() 把只含一个元素的张量转成普通的 Python 整数
- print(f"Predicted class: {predicted_class}") # 打印预测结果，例如 "Predicted class: 0"
- # 下面三行是把送入模型的图片显示出来，方便人工核对预处理结果
- img_display = input_tensor.squeeze().cpu().numpy() # 把张量变回 NumPy 数组以便转成图片：
- # squeeze() 去掉所有大小为 1 的维度，[1,1,28,28] -> [28,28]
- # .cpu() 把张量从 GPU 移回 CPU（NumPy 只能处理 CPU 上的数据）
- # .numpy() 把 PyTorch 张量转成 NumPy 数组，此时像素值仍是 0.0~1.0 的浮点数
- image = Image.fromarray((img_display * 255).astype(np.uint8)) # 把 NumPy 数组转回 PIL 图片：
- # *255 把像素值从 0~1 还原回 0~255
- # astype(np.uint8) 转成 8 位无符号整数（图片标准格式）
- # Image.fromarray 根据数组创建灰度图对象
- image.show() # 调用系统默认的图片查看器显示这张图片
|
