模型训练好后预测 numpy图片

用户6021899

发布于 2022-11-18 14:10:12

5180

文章被收录于专栏：Python编程 pyqt matplotlibPython编程 pyqt matplotlib

神经网络训练好后，预测若干张图片（实际上是numpy 数组），可将numpy 数组转换成 size 为 (batch,channels, height, width), 类型为float 的 tersor 后，直接预测。不必有实体图片文件，不必建数据集。

注意要手动将数据归一化，mean值和标准差与训练集的 transforms 的归一化参数一致。

if __name__ == "__main__":
    model_path = "diecode_mode.pth"

    #device = "cuda" if torch.cuda.is_available  else "cpu"  # Get cpu or gpu device for training.
    device ="cpu"
    net = NeuralNetwork() # .to(device)
    net.load_state_dict(torch.load(model_path))

    images = np.random.random((4,1,50,36)) # batch, channels, height, width
    mean, std = 0.5, 0.5
    images = (images-mean) / std  # 对应 transforms.Normalize()

    # .to(torch.float) ：权重默认都是浮点数，所以input也必须是浮点数
    # 否则报错 RuntimeError: expected scalar type Double but found Float
    images = torch.from_numpy(images).to(torch.float)

    #print(images)
    with torch.no_grad():
        outputs = net(images)
        _, predictions = torch.max(outputs, 1)
        print(predictions)
        print([classes[x] for x in predictions])