Пример #1
0
def main():
    """Run Prediction"""

    # Measures total program runtime by collecting start time
    start_time = time.time()

    # get input args
    in_arg = get_input_args()

    # Function that checks command line arguments using in_arg
    print_command_line_arguments(in_arg)
    print()

    # get dataloader to load labels and process image
    dataloader = ModelDataLoader('data')
    category_names = dataloader.get_label_dict(in_arg.category_names)

    # use modelloader to load checkpoint
    loader = ModelLoader(in_arg.checkpoint_dir)

    if loader.checkpoint_exists(in_arg.checkpoint):
        print("Loading checkpoint %s" %
              (loader.get_checkpoint_path(in_arg.checkpoint)))
        model = loader.load_checkpoint(
            in_arg.checkpoint, model_use_gpu=in_arg.gpu, with_trainer=False)
    else:
        print("Checkpoint '%s' does not exist. Exiting." %
              (loader.get_checkpoint_path(in_arg.checkpoint)))
        return

    print()

    # convert image
    image = dataloader.process_image(in_arg.image)

    # run prediction
    probs, classes = model.predict(image, in_arg.top_k)
    labels = [category_names[c] for c in classes]

    for prob, clas, label in zip(probs, classes, labels):
        print("%5.3f%% %s [%s]" % (prob*100, label, clas))

    print()
    print("Total Prediction Duration: %.3fs" % (time.time() - start_time))
Пример #2
0
def main():
    """Run Training Session"""

    # Measures total program runtime by collecting start time
    start_time = time.time()

    # get input args
    in_arg = get_input_args()

    # Function that checks command line arguments using in_arg
    print_command_line_arguments(in_arg)
    print()

    # load datasets
    dataloader = ModelDataLoader(in_arg.data_dir)
    train_dataset, train_dataloader = dataloader.get_train_data()
    valid_dataset, valid_dataloader = dataloader.get_validation_data()
    test_dataset, test_dataloader = dataloader.get_test_data()

    # Use model loader to load existing checkpoint
    loader = ModelLoader(in_arg.checkpoint_dir)

    if loader.checkpoint_exists(in_arg.checkpoint):
        # load checkpoint
        print("Loading checkpoint %s" %
              (loader.get_checkpoint_path(in_arg.checkpoint)))
        model, trainer = loader.load_checkpoint(
            in_arg.checkpoint, model_use_gpu=in_arg.gpu)
        print("Epochs trained so far: %d" % (trainer.trained_epochs))
    else:
        # no checkpoint, create fresh model using input arguments
        print("Checkpoint '%s' does not exist" %
              (loader.get_checkpoint_path(in_arg.checkpoint)))
        model = Model(train_dataset.class_to_idx,
                      arch=in_arg.arch, use_gpu=in_arg.gpu,
                      hidden_units_1=in_arg.hidden_units_1, hidden_units_2=in_arg.hidden_units_2,
                      dropout_1=in_arg.dropout_1, dropout_2=in_arg.dropout_2)
        trainer = ModelTrainer(model, learning_rate=in_arg.learning_rate)

    print()
    print("Model training in session...")
    print()

    epochs = in_arg.epochs

    # train model and print results
    for result in trainer.train_epochs(epochs, train_dataloader, valid_dataloader):
        print(
            "Epoch: %3d/%3d" % (result['epoch']+1, epochs),
            " | Train Loss: %10.5f" % (result['train_loss']),
            " | Validation Loss: %10.5f" % (result['validation_loss']),
            " | Validation Acc: %6.3f%%" % (
                result['validation_accuracy'] * 100),
            " | Duration: %10.3fs" % (result['duration'])
        )

    print()
    print("Testing model against test data...")
    print()

    # test model against test data
    test_result = trainer.test(test_dataloader)

    print(
        "Test Loss: %10.5f" % (test_result['test_loss']),
        " | Test Acc: %6.3f%%" % (test_result['test_accuracy'] * 100),
        " | Duration: %10.3fs" % (test_result['duration'])
    )

    # save checkpoint
    loader.save_checkpoint(in_arg.checkpoint, model, trainer)

    print()
    print("Total Train Duration: %.3fs" % (time.time() - start_time))