File size: 3,237 Bytes
4b7c478 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | import os
import torch
import numpy as np
from torch import nn
import torch.nn.functional as F
from tqdm import tqdm
from torchvision import datasets
import torchvision.transforms as T
from .helpers import get_data_location
class Predictor(nn.Module):
def __init__(self, model, class_names, mean, std):
super().__init__()
self.model = model.eval()
self.class_names = class_names
# We use nn.Sequential and not nn.Compose because the former
# is compatible with torch.script, while the latter isn't
self.transforms = nn.Sequential(
T.Resize([256, ]), # We use single int value inside a list due to torchscript type restrictions
T.CenterCrop(224),
T.ConvertImageDtype(torch.float),
T.Normalize(mean.tolist(), std.tolist())
)
self.act = nn.Softmax(dim=1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
with torch.no_grad():
# 1. apply transforms
x = self.transforms(x)
# 2. get the logits
x = self.model(x)
# 3. apply softmax
# HINT: remmeber to apply softmax across dim=1
x = self.act(x)
return x
def predictor_test(test_dataloader, model_reloaded):
"""
Test the predictor. Since the predictor does not operate on the same tensors
as the non-wrapped model, we need a specific test function (can't use one_epoch_test)
"""
folder = get_data_location()
test_data = datasets.ImageFolder(os.path.join(folder, "test"), transform=T.ToTensor())
pred = []
truth = []
for x in tqdm(test_data, total=len(test_dataloader.dataset), leave=True, ncols=80):
softmax = model_reloaded(x[0].unsqueeze(dim=0))
idx = softmax.squeeze().argmax()
pred.append(int(x[1]))
truth.append(int(idx))
pred = np.array(pred)
truth = np.array(truth)
print(f"Accuracy: {(pred==truth).sum() / pred.shape[0]}")
return truth, pred
######################################################################################
# TESTS
######################################################################################
import pytest
@pytest.fixture(scope="session")
def data_loaders():
from .data import get_data_loaders
return get_data_loaders(batch_size=2)
def test_model_construction(data_loaders):
from .model import MyModel
from .helpers import compute_mean_and_std
mean, std = compute_mean_and_std()
model = MyModel(num_classes=3, dropout=0.3)
dataiter = iter(data_loaders["train"])
images, labels = next(dataiter)
predictor = Predictor(model, class_names=['a', 'b', 'c'], mean=mean, std=std)
out = predictor(images)
assert isinstance(
out, torch.Tensor
), "The output of the .forward method should be a Tensor of size ([batch_size], [n_classes])"
assert out.shape == torch.Size(
[2, 3]
), f"Expected an output tensor of size (2, 3), got {out.shape}"
assert torch.isclose(
out[0].sum(),
torch.Tensor([1]).squeeze()
), "The output of the .forward method should be a softmax vector with sum = 1"
|