Single image prediction for a trained UNET model gone wrong

vision
1

I have trained my UNet model and got a good accuracy result though for a single image prediction I am suffering more and more. I am using google colab for training purposes and saving my model into google drive. After loading the trained model prediction is not good, say only about 30/40 % though I was getting 90% accuracy during training. What might be the possible cause for this?
My training loop:

for batch in train_loader:
    model.train()
    images, labels = batch
    labels = labels.squeeze(dim=1)
    labels = labels.to(torch.long)

    iteration += 1

    preds = model(images)
    loss = criterion(preds, labels)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    predicted = torch.argmax(preds.data, 1)
    total_train += labels.nelement()
    correct_train += predicted.eq(labels.data).sum().item()
    train_accuracy = 100 * correct_train / total_train

from PIL import Image
from torchvision import transforms
input_image = Image.open('drive/My Drive/BackgroundImageRemoval/Test Images/400pixel.jpg')
preprocess = transforms.Compose([
    transforms.CenterCrop(248),
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

input_tensor = preprocess(input_image)

input_batch = input_tensor.unsqueeze(0)
with torch.no_grad():
  output2 = model(input_batch)
  output2 = output2.squeeze(0)
output2_predictions = output2.argmax(0)
r = Image.fromarray(output2_predictions.byte().cpu().numpy()).resize(input_image.size)
r
2

It’s hard to tell, where the accuracy drop might be coming from and I would recommend to load a specific image sample during training and validation and compare the output of these runs.
If your model doesn’t give the same output (in eval() mode), than something might be different in the data loading or preprocessing.

1 Like