lama

1
#!/usr/bin/env python3
2

3

4
import os
5
from argparse import ArgumentParser
6

7

8
def ssim_fid100_f1(metrics, fid_scale=100):
9
    ssim = metrics.loc['total', 'ssim']['mean']
10
    fid = metrics.loc['total', 'fid']['mean']
11
    fid_rel = max(0, fid_scale - fid) / fid_scale
12
    f1 = 2 * ssim * fid_rel / (ssim + fid_rel + 1e-3)
13
    return f1
14

15

16
def find_best_checkpoint(model_list, models_dir):
17
    with open(model_list) as f:
18
        models = [m.strip() for m in f.readlines()]
19
    with open(f'{model_list}_best', 'w') as f:
20
        for model in models:
21
            print(model)
22
            best_f1 = 0
23
            best_epoch = 0
24
            best_step = 0
25
            with open(os.path.join(models_dir, model, 'train.log')) as fm:
26
                lines = fm.readlines()
27
                for line_index in range(len(lines)):
28
                    line = lines[line_index]
29
                    if 'Validation metrics after epoch' in line:
30
                        sharp_index = line.index('#')
31
                        cur_ep = line[sharp_index + 1:]
32
                        comma_index = cur_ep.index(',')
33
                        cur_ep = int(cur_ep[:comma_index])
34
                        total_index = line.index('total ')
35
                        step = int(line[total_index:].split()[1].strip())
36
                        total_line = lines[line_index + 5]
37
                        if not total_line.startswith('total'):
38
                            continue
39
                        words = total_line.strip().split()
40
                        f1 = float(words[-1])
41
                        print(f'\tEpoch: {cur_ep}, f1={f1}')
42
                        if f1 > best_f1:
43
                            best_f1 = f1
44
                            best_epoch = cur_ep
45
                            best_step = step
46
            f.write(f'{model}\t{best_epoch}\t{best_step}\t{best_f1}\n')
47

48

49
if __name__ == '__main__':
50
    parser = ArgumentParser()
51
    parser.add_argument('model_list')
52
    parser.add_argument('models_dir')
53
    args = parser.parse_args()
54
    find_best_checkpoint(args.model_list, args.models_dir)
55