模型训练好后 预测 numpy图片

2022-11-18 14:10:12 浏览数 (1)

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

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

代码语言:javascript复制
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])

0 人点赞