PyTorch 图像分类模型测试代码示例

import torch
import torch.nn.functional as F
import torch.nn as nn

# 创建网络实例
class MyNetwork(nn.Module):
    def __init__(self):
        super(MyNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.fc = nn.Linear(339 * 256, 4)

    def forward(self, x):
        x = self.flatten(x)
        x = self.fc(x)
        return x

# 加载保存的网络参数
network = MyNetwork()
network.load_state_dict(torch.load('final_model.pt'))

# 读取标签数据
file_path = 'biaoqian.txt'
tensor_list = []
with open(file_path, 'r') as file:
    lines = file.readlines()
    for line in lines:
        line = line.strip()
        numbers = line.split()
        tensor = torch.tensor([float(num) for num in numbers])
        tensor_list.append(tensor)

# 读取测试数据
pt_file_path = 'test_data.pt'
test_data = torch.load(pt_file_path)

# 将测试数据转换为张量
test_tensors = [torch.tensor(tensor) for tensor in test_data]

# 在测试集上进行预测
network.eval()
with torch.no_grad():
    predictions = []
    for test_input_tensor in test_tensors:
        output = network(test_input_tensor.unsqueeze(0))
        prediction = torch.argmax(output)
        predictions.append(prediction.item())

# 打印预测结果
print(predictions)

这段代码演示了如何使用PyTorch加载训练好的图像分类模型并进行预测。

代码解释：

定义网络结构: 代码首先定义了名为 MyNetwork 的网络结构，这部分需要与你训练的模型结构一致。
加载模型参数: 使用 torch.load() 加载训练好的模型参数 (.pt 文件)。
读取标签数据: 从 biaoqian.txt 文件中读取标签数据，并将其转换为张量形式。
读取测试数据: 从 test_data.pt 文件中读取测试数据。
进行预测: 使用 network.eval() 将模型设置为评估模式，然后对测试数据进行预测，并将预测结果存储在 predictions 列表中。
打印结果: 最后，打印预测结果。

注意事项：

将代码中的文件名替换为你实际使用的文件名。
确保你的测试数据与训练数据格式一致。
可以根据需要修改代码以适应你的具体需求，例如将预测结果保存到文件或进行可视化。