如何使用 PyTorch 模型进行预测？

Question

我创建了一个PyTorch模型来分类图像。 我通过state_dict一次保存了整个模型，如下：

torch.save(model.state_dict(), "model1_statedict")
torch.save(model, "model1_complete")

我该如何使用这些模型？我想用一些图像来检查它们的表现是否良好。

我正在使用以下代码加载模型：

model = torch.load(path_model)
model.eval()

这个可以，但我不知道如何使用它来预测新图片。

- lxg95

---

我修改了你的问题，因为在这里询问资源，比如教程，是不被允许的。 - Nicolas Gervais

好的，抱歉，我不知道，感谢您的编辑。 - lxg95

不允许？愚蠢 - france1

Answer 1

def predict(self, test_images):
    self.eval()
    # model is self(VGG class's object)
    
    count = test_images.shape[0]
    result_np = []
        
    for idx in range(0, count):
        # print(idx)
        img = test_images[idx, :, :, :]
        img = np.expand_dims(img, axis=0)
        img = torch.Tensor(img).permute(0, 3, 1, 2).to(device)
        # print(img.shape)
        pred = self(img)
        pred_np = pred.cpu().detach().numpy()
        for elem in pred_np:
            result_np.append(elem)
    return result_np

网络结构是VGG-19，参考我的源代码。

像这样的架构：

class VGG(object):
    def __init__(self):
    ...


    def train(self, train_images, valid_images):
        train_dataset = torch.utils.data.Dataset(train_images)
        valid_dataset = torch.utils.data.Dataset(valid_images)

        trainloader = torch.utils.data.DataLoader(train_dataset)
        validloader = torch.utils.data.DataLoader(valid_dataset)

        self.optimizer = Adam(...)
        self.criterion = CrossEntropyLoss(...)
    
        for epoch in range(0, epochs):
            ...
            self.evaluate(validloader, model=self, criterion=self.criterion)
    ...

    def evaluate(self, dataloader, model, criterion):
        model.eval()
        for i, sample in enumerate(dataloader):
    ...

    def predict(self, test_images):
    
    ...

if __name__ == "__main__":
    network = VGG()
    trainset, validset = get_dataset()    # abstract function for showing
    testset = get_test_dataset()
    
    network.train(trainset, validset)

    result = network.predict(testset)

- bolero

Answer 2

PyTorch模型是一个函数。你提供适当定义的输入，它返回一个输出。如果你只想在给定特定输入图像的情况下视觉检查输出，只需调用它：

model.eval()
output = model(example_image)

- iacob