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,)