Pytorch基础知识（4）模型部署

为了部署模型，我们需要使用前一节中描述的方法加载模型。一旦将模型加载到内存中，我们就可以将新数据传递给模型。

1. 获取验证集上的一个样本tensor

n = 100
x = x_val[n]
y = y_val[n]
print(x.shape)
plt.imshow(x.numpy()[0], cmap="gray")

# torch.Size([1, 28, 28])

2. 然后，预处理tensor

# we use unsqueeze to expand dimensions to 1*C*H*W
x = x.unsqueeze(0)

# convert to torch.float32
x = x.type(torch.float)

# move to cuda device
x = x.to(device)

3. 获得模型预测结果

# get model output
output = _model(x)

# get predicted class
pred = output.argmax(dim=1, keepdim=True)
print(pred.item(), y.item())

# 6 6