test_batch_size.py

"""
Test whether modifying the batch size changes the result

Authors: Marijn van Vliet <[email protected]>
"""
import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np

import hmax

# Initialize the model with the universal patch set
print('Constructing model')
model = hmax.HMAX('./universal_patch_set.mat')

# A folder with example images
example_images = datasets.ImageFolder(
    './example_images/',
    transform=transforms.Compose([
        transforms.Grayscale(),
        transforms.ToTensor(),
        transforms.Lambda(lambda x: x * 255),
    ])
)

# A dataloader that will run through all example images in one batch
dataloader = DataLoader(example_images, batch_size=10)

# Determine whether there is a compatible GPU available
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Run the model on the example images
print('Running model on', device)
model = model.to(device)
for X, y in dataloader:
    s1_1, c1_1, s2_1, c2_1 = model.get_all_layers(X.to(device))

# Run the model again, but with a smaller batch size
dataloader = DataLoader(example_images, batch_size=2)
for X, y in dataloader:
    s1_2, c1_2, s2_2, c2_2 = model.get_all_layers(X.to(device))

# Check equivalency of output for the last batch
for scale in range(8):
    assert np.array_equal(s1_1[scale][-2:], s1_2[scale])
    assert np.array_equal(c1_1[scale][-2:], c1_2[scale])
    for scale2 in range(8):
        assert np.array_equal(s2_1[scale][scale2][-2:], s2_2[scale][scale2])
    assert np.array_equal(c2_1[scale][-2:], c2_2[scale])