Views
No views yet
1
2## Approach 1: use pytorch to predict
3import torch
4from model import CIFARCNN
5
6# Evaluate model checkpoints
7model = CIFARCNN.load_from_checkpoint("model.ckpt")
8model.eval()
9x = torch.randn(4, 3, 32, 32).to(model.device)
10
11with torch.no_grad():
12 predictions = model(x) # the lightning module should implement forward func
13print(predictions.shape) # should be [4, 10]1import torch
2from model import CIFARCNN
3from lightning import Trainer
4
5test_dataloader = DataLoader(...)
6model = CIFARCNN.load_from_checkpoint("model.ckpt") # lightning will move model to default device
7trainer = Trainer()
8
9trainer.test(model, test_dataloader)1import matplotlib.pyplot as plt
2
3cifar10_labels = {
4 0: "airplane",
5 1: "automobile",
6 2: "bird",
7 3: "cat",
8 4: "deer",
9 5: "dog",
10 6: "frog",
11 7: "horse",
12 8: "ship",
13 9: "truck",
14}
15
16samples, labels = next(iter(train_loader))
17predicts = trainer.predict(model, samples)
18labels = predicts.argmax(dim=1)
19
20fig, axes = plt.subplots(2, 5, figsize=(10, 4))
21for i, ax in enumerate(axes.flatten()):
22 ax.imshow(samples[i].permute(1, 2, 0))
23 ax.set_title(f"{cifar10_labels[labels[i].item()]}")
24 ax.axis("off")
25plt.show()