加载state_dict 预测单张图片

以下是使用PyTorch对单张图片进行模型预测的示例代码，这里使用`state_dict`加载模型：

import torch
import torchvision.transforms as transforms
from PIL import Image

# 定义模型结构
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = torch.nn.Conv2d(10, 20, kernel_size=5)
        self.fc1 = torch.nn.Linear(320, 50)
        self.fc2 = torch.nn.Linear(50, 10)

    def forward(self, x):
        x = torch.nn.functional.relu(torch.nn.functional.max_pool2d(self.conv1(x), 2))
        x = torch.nn.functional.relu(torch.nn.functional.max_pool2d(self.conv2(x), 2))
        x = x.view(-1, 320)
        x = torch.nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return torch.nn.functional.log_softmax(x, dim=1)

# 加载模型参数
model = Net()
model.load_state_dict(torch.load('mnist_model_state_dict.pt'))

# 定义变换
transform = transforms.Compose([
    transforms.Resize((28,28)),
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载图片
image = Image.open('test.png').convert('L')
image = transform(image)

# 进行预测
with torch.no_grad():
    output = model(image.unsqueeze(0))
    _, predicted = torch.max(output.data, 1)
    print('Predicted:', predicted.item())

其中，`mnist_model_state_dict.pt`是你已经训练好的模型参数文件，需要根据你的实际情况进行修改。在定义模型结构时，需要与训练时保持一致。在进行预测前，需要对图片进行变换，将其转化为模型所需的格式。最后，使用`model`对变换后的图片进行预测，输出预测结果。