experiment_launcher/examples/test.py at master · im-Kitsch/experiment_launcher · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
import argparse
import os
import time

from experiment_launcher import get_default_params, run_experiment, add_launcher_base_args, save_args, save_success


def train(a=1,
          b_c=1,
          boolean=True,
          default='dft',
          seed=0,  # This argument is mandatory
          results_dir='/tmp'  # This argument is mandatory
          ):
    filename = os.path.join(results_dir, 'log_' + str(seed) + '.txt')
    os.makedirs(results_dir, exist_ok=True)
    print('Running experiment with seed', str(seed),
          'and parameters a =', str(a),
          'b_c =', str(b_c),
          'boolean =', str(boolean),
          'default =', default)
    with open(filename, 'w') as file:
        file.write('Default parameters:\n')
        file.write('seed: ' + str(seed) + '\n')
        file.write('results_dir: ' + str(results_dir) + '\n')

        file.write('Experiment parameters:\n')
        file.write('a: ' + str(a) + '\n')
        file.write('b-c: ' + str(b_c) + '\n')
        file.write('boolean: ' + str(boolean) + '\n')
        file.write('default: ' + default + '\n')
    return


def experiment(a=1,
               b_c=1,
               boolean=True,
               default='dft',
               seed=0,  # This argument is mandatory
               results_dir='/tmp'  # This argument is mandatory
               ):

    ####################################################################################################################
    # SETUP
    # Create results directory
    results_dir = os.path.join(results_dir, str(seed))
    os.makedirs(results_dir, exist_ok=True)
    # Save arguments
    save_args(results_dir, locals(), git_repo_path='./')
    save_success(results_dir, locals(), False)

    ####################################################################################################################
    # EXPERIMENT

    train(a,  b_c, boolean,  default,  seed,  results_dir)

    save_success(results_dir, locals(), True)

def parse_args():
    parser = argparse.ArgumentParser()

    # Place your experiment arguments here
    arg_test = parser.add_argument_group('Test')
    arg_test.add_argument("--a", type=int)
    arg_test.add_argument("--b-c", type=int)
    arg_test.add_argument("--boolean", action='store_true')
    arg_test.add_argument('--default', type=str)

    # Leave unchanged
    parser = add_launcher_base_args(parser)
    parser.set_defaults(**get_default_params(experiment))
    args = parser.parse_args()
    return vars(args)


if __name__ == '__main__':
    # Leave unchanged
    args = parse_args()
    run_experiment(experiment, args)