Hackathon-Attribution/test_sensitivity.py at main · pinplex/Hackathon-Attribution · 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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
from argparse import ArgumentParser, Namespace
import torch
import os
import shutil
from tqdm import tqdm

from hackathon.model_runner import ModelRunner
from hackathon.base_model import BaseModel
from hackathon.models.attn import MultiheadAttn as attn_model
from hackathon.models.Conv1D import Conv1D as conv1d_model
from hackathon.models.linear import Linear as linear_model
from hackathon.models.LSTM import LSTM as lstm_model
from hackathon.models.multimodel import EfficiencyModel as efficiency_model
from hackathon.models.resnet import ResNetModule as resnet_model
from hackathon.models.simplemlp import SimpleMLP as simplemlp_model
from hackathon.models.attn_nores import MultiheadAttnNoRes as attn_nores_model

model_funs = [ # uncomment models you like to process
    #attn_model,
    #conv1d_model,
    #linear_model,
    #lstm_model,
    #simplemlp_model,
    attn_nores_model
]


def main(args: Namespace):
    model: BaseModel

    print('\n+++ Predictions are saved to `./hackathon/logs/<model_name>/sensitivity/<CO2|causal>/predictions.nc>`+++\n')

    for model in (pbar0 := tqdm(model_funs)):
        model_name = model.__module__.split('.')[-1].lower()
        log_dir = f'./hackathon/logs/{model_name}/sensitivity'
        if os.path.isdir(log_dir):
            shutil.rmtree(log_dir)
        checkpoint_path = f'./hackathon/logs/{model_name}/xval/final/final.ckpt'

        runner_CO2 = ModelRunner(
            log_dir=log_dir,
            quickrun=args.quickrun,
            seed=910,
            data_dir='./simple_gpp_model/data/CMIP6/predictor-variables_historical+ssp585+GPP_no-CO2-change.nc')
        runner_causal = ModelRunner(
            log_dir=log_dir,
            quickrun=args.quickrun,
            seed=910,
            data_dir='./simple_gpp_model/data/CMIP6/predictor-variables_historical+ssp585+GPP_non-causal-constant.nc')

        model = torch.load(checkpoint_path)

        pbar0.set_description(f'Model     {"<"+model_name+">":>30}')

        for runner, runner_name in (pbar1 := tqdm(zip(
                [runner_CO2, runner_causal],
                ['CO2', 'causal']), leave=False)):

            pbar1.set_description(f'Predicting {"<"+runner_name+">":>30}')

            if runner_name == 'CO2':
                data_module = runner.data_setup(
                    fold=-1,
                    custom_test_sel={
                        'time': slice('1850', '2100'),
                        'location': range(1, 2)
                    },
                    batch_size=1
                )
            else:
                data_module = runner.data_setup(
                    fold=-1,
                    custom_test_sel={
                        'time': slice('1850', '2100'),
                        'location': range(1, 2)
                    },
                    batch_size=1
                )

            trainer = runner.trainer_setup(
                version=runner_name,
                accelerator=None if args.gpu == -1 else 'gpu',
                devices=None if args.gpu == -1 else f'{args.gpu},'
            )

            runner.predict(
                trainer=trainer,
                model=model,
                datamodule=data_module,
                version=runner_name
            )


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument(
        '--quickrun',
        action='store_true',
        help='Quick development run, explanations are computed for only one location .')
    parser.add_argument(
        '--gpu',
        type=int,
        default=-1,
        help='GPU ID to use, -1 (default) deactivates GPU.')

    args = parser.parse_args()

    main(args)