def test_pytorch_gtsrb():
classifier_module = import_module(
"armory.baseline_models.pytorch.micronnet_gtsrb")
classifier_fn = getattr(classifier_module, "get_art_model")
preprocessing_fn = getattr(classifier_module, "preprocessing_fn")
classifier = classifier_fn(model_kwargs={}, wrapper_kwargs={})
train_dataset = datasets.german_traffic_sign(
split="train",
epochs=5,
batch_size=128,
dataset_dir=DATASET_DIR,
preprocessing_fn=preprocessing_fn,
)
test_dataset = datasets.german_traffic_sign(
split="test",
epochs=1,
batch_size=128,
dataset_dir=DATASET_DIR,
preprocessing_fn=preprocessing_fn,
)
classifier.fit_generator(
train_dataset,
nb_epochs=5,
)
accuracy = 0
for _ in range(test_dataset.batches_per_epoch):
x, y = test_dataset.get_batch()
predictions = classifier.predict(x)
accuracy += np.sum(np.argmax(predictions, axis=1) == y) / len(y)
assert (accuracy / test_dataset.batches_per_epoch) > 0.8
def test_german_traffic_sign():
for split, size in [("train", 39209), ("test", 12630)]:
batch_size = 1
epochs = 1
dataset = datasets.german_traffic_sign(
split=split, epochs=epochs, batch_size=batch_size, dataset_dir=DATASET_DIR,
)
assert dataset.size == size
x, y = dataset.get_batch()
# sign image shape is variable so we don't compare 2nd dim
assert x.shape[:1] + x.shape[3:] == (batch_size, 3)
assert y.shape == (batch_size,)