Function main

utils/plot-results.py:154–191

(args)

def main(args):
    os.makedirs(args.out_dir, exist_ok=True)

    # setup plots
    sns.set_style('darkgrid')
    sns.set()

    result_files = glob.glob(os.path.join(args.results_dir, "*.json"), recursive=True)
    for resfile in result_files:
        with open(resfile, 'r') as infile:
            raw_data = json.load(infile)
        data = pd.DataFrame(raw_data.values())
        dataname = data.dataset[0]
        if args.dataset != "all" and dataname != args.dataset:
             continue

        bn_data = data[data.pretrain_iters.str.contains("bn")]
        data = data[~data.pretrain_iters.str.contains("bn")]        
        data.pretrain_iters = pd.to_numeric(data.pretrain_iters, errors='coerce')
        # convert batchnorm iters to a new type of data

        try:
            data = data.astype({
                "pretrain_iters": int
            })
        # pokemon exceptions!
        except Exception as exp:
            print(f"WARNING: Unable to parse pretrain_iters as int for {dataname}")
        if data.result.max() < 1:
            data.result *= 100
        data.basetrain = data.basetrain.replace("imagenet_r50_supervised", "supervised imagenet init")
        data.basetrain = data.basetrain.replace("moco_v2_800ep", "moco imagenet init")
        data.basetrain = data.basetrain.replace("no", "random init")
        data.basetrain = data.basetrain.replace("none", "random init")        
        gen_plots(data, bn_data, {
            'out_dir': args.out_dir,
            'data_name': dataname
        })

if __name__ == "__main__":
    main(parse_args())