train_and_eval_autokeras.py

import argparse
import os
import time

import autokeras as ak
import kerastuner
import medmnist
import numpy as np
import tensorflow as tf
from medmnist import INFO, Evaluator
from medmnist.info import DEFAULT_ROOT
from tensorflow.keras.models import load_model


def main(data_flag, num_trials, input_root, output_root, gpu_ids, run, model_path):

    os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" 
    os.environ["CUDA_VISIBLE_DEVICES"]=str(gpu_ids)

    info = INFO[data_flag]
    task = info['task']
    _ = getattr(medmnist, INFO[data_flag]['python_class'])(
            split="train", root=input_root, download=True)

    output_root = os.path.join(output_root, data_flag, time.strftime("%y%m%d_%H%M%S"))
    if not os.path.isdir(output_root):
        os.makedirs(output_root)

    npz_file = np.load(os.path.join(input_root, "{}.npz".format(data_flag)))

    x_train = npz_file['train_images']
    y_train = npz_file['train_labels']
    x_val = npz_file['val_images']
    y_val = npz_file['val_labels']
    x_test = npz_file['test_images']
    y_test = npz_file['test_labels']

    if model_path is not None:
        model = load_model(model_path, custom_objects=ak.CUSTOM_OBJECTS)
        test(model, data_flag, x_train, 'train', output_root, run)
        test(model, data_flag, x_val, 'val', output_root, run)
        test(model, data_flag, x_test, 'test', output_root, run)
        
    if num_trials == 0:
        return

    model = train(data_flag, x_train, y_train, x_val, y_val, num_trials, output_root, run)

    test(model, data_flag, x_train, 'train', output_root, run)
    test(model, data_flag, x_val, 'val', output_root, run)
    test(model, data_flag, x_test, 'test', output_root, run)


def train(data_flag, x_train, y_train, x_val, y_val, num_trials, output_root, run):

    clf = ak.ImageClassifier(
        project_name=data_flag,
        distribution_strategy=tf.distribute.MirroredStrategy(),
        metrics=['AUC', 'accuracy'],
        objective=kerastuner.Objective("val_auc", direction="max"),
        overwrite=True,
        max_trials=num_trials
    )

    clf.fit(
        x_train,
        y_train,
        validation_data=(x_val, y_val),
        epochs=20
    )

    model = clf.export_model()

    try:
        model.save(os.path.join(output_root, '%s_autokeras_%s' % (data_flag, run)), save_format="tf")
    except Exception:
        model.save(os.path.join(output_root, '%s_autokeras_%s.h5' % (data_flag, run)))

    return model


def test(model, data_flag, x, split, output_root, run):

    evaluator = medmnist.Evaluator(data_flag, split)
    y_score = model.predict(x)
    auc, acc = evaluator.evaluate(y_score, output_root, run)
    print('%s  auc: %.5f  acc: %.5f' % (split, auc, acc))

    return auc, acc


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--data_flag',
                        default='organmnist3d',
                        type=str)
    parser.add_argument('--input_root',
                        default=DEFAULT_ROOT,
                        type=str)
    parser.add_argument('--output_root',
                        default='./autokeras',
                        type=str)
    parser.add_argument('--gpu_ids',
                        default='0',
                        type=str)
    parser.add_argument('--run',
                        default='model1',
                        help='to name a standard evaluation csv file, named as {flag}_{split}_[AUC]{auc:.3f}_[ACC]{acc:.3f}@{run}.csv',
                        type=str)
    parser.add_argument('--model_path',
                        default=None,
                        help='root of the pretrained model to test',
                        type=str)
    parser.add_argument('--num_trials',
                        default=20,
                        help='max_trials of autokeras search space, the script would only test model if num_trials=0',
                        type=int)

    args = parser.parse_args()
    data_flag = args.data_flag
    input_root = args.input_root
    output_root = args.output_root
    gpu_ids = args.gpu_ids
    run = args.run
    model_path = args.model_path
    num_trials = args.num_trials

    main(data_flag, num_trials, input_root, output_root, gpu_ids, run, model_path)