hyperparam_search.py 文件源码

def main(training_label_prefix,
         dataset_name=None,
         epochs=None,
         time_limit=None,
         num_gpus=None):
    epochs = int(epochs) if epochs else None
    time_limit = parse_timedelta(time_limit) if time_limit else None
    num_gpus = int(num_gpus) if num_gpus else None
    search = HyperparamSearch(training_label_prefix=training_label_prefix,
                              dataset_name=dataset_name,
                              epochs=epochs,
                              time_limit=time_limit,
                              num_gpus=num_gpus)

    def handler(signum, frame):
        logging('Stopping hyperparam search..')
        with search.lock:
            search.stop()
            for index, running_command in search.running_commands:
                try:
                    label = search.training_label(index)
                    logging('Sending SIGINT to {}..'.format(label))
                    running_command.signal(signal.SIGINT)
                except OSError:  # The process might have exited before
                    logging('{} might have terminated before.'.format(label))
                except:
                    traceback.print_exc(file=sys.stderr)
            logging('All training processes have been sent SIGINT.')
    signal.signal(signal.SIGINT, handler)

    # We need to execute search.run() in another thread in order for Semaphore
    # inside it doesn't block the signal handler. Otherwise, the signal handler
    # will be executed after any training process finishes the whole epoch.

    executor = ThreadPoolExecutor(max_workers=1)
    executor.submit(search.run)
    # wait must be True in order for the mock works,
    # see the unit test for more details
    executor.shutdown(wait=True)