-
Notifications
You must be signed in to change notification settings - Fork 43
/
Copy pathtest_finetune.py
72 lines (50 loc) · 2.56 KB
/
test_finetune.py
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
import os
import shutil
import gc
import pytest
import timm
import torch
from terratorch.cli_tools import build_lightning_cli
from terratorch.registry import BACKBONE_REGISTRY
@pytest.fixture(autouse=True)
def setup_and_cleanup(model_name):
model_instance = BACKBONE_REGISTRY.build(model_name)
state_dict = model_instance.state_dict()
torch.save(state_dict, os.path.join("tests", model_name + ".pt"))
yield # everything after this runs after each test
os.remove(os.path.join("tests", model_name + ".pt"))
if os.path.isdir(os.path.join("tests", "all_ecos_random")):
shutil.rmtree(os.path.join("tests", "all_ecos_random"))
@pytest.mark.parametrize("model_name", ["prithvi_eo_v1_100", "prithvi_eo_v2_300", "prithvi_swin_B", "prithvi_swin_L", "prithvi_eo_v2_600"])
@pytest.mark.parametrize("case", ["fit", "test", "validate", "compute_statistics"])
def test_finetune_multiple_backbones(model_name, case):
command_list = [case, "-c", f"tests/resources/configs/manufactured-finetune_{model_name}.yaml"]
_ = build_lightning_cli(command_list)
gc.collect()
@pytest.mark.parametrize("model_name", ["prithvi_swin_B"])
@pytest.mark.parametrize("case", ["fit", "test", "validate"])
def test_finetune_bands_intervals(model_name, case):
command_list = [case, "-c", f"tests/resources/configs/manufactured-finetune_{model_name}_band_interval.yaml"]
_ = build_lightning_cli(command_list)
gc.collect()
@pytest.mark.parametrize("model_name", ["prithvi_swin_B"])
@pytest.mark.parametrize("case", ["fit", "test", "validate"])
def test_finetune_bands_str(model_name, case):
command_list = [case, "-c", f"tests/resources/configs/manufactured-finetune_{model_name}_string.yaml"]
_ = build_lightning_cli(command_list)
gc.collect()
@pytest.mark.parametrize("model_name", ["prithvi_eo_v2_300"])
@pytest.mark.parametrize("case", ["fit", "test", "validate"])
def test_finetune_pad(case):
command_list = [case, "-c", f"tests/resources/configs/manufactured-finetune_prithvi_pixelwise_pad.yaml"]
_ = build_lightning_cli(command_list)
gc.collect()
@pytest.mark.parametrize("model_name", ["prithvi_swin_B"])
def test_finetune_metrics_from_file(model_name):
model_instance = timm.create_model(model_name)
state_dict = model_instance.state_dict()
torch.save(state_dict, os.path.join("tests/resources/configs/", model_name + ".pt"))
# Running the terratorch CLI
command_list = ["fit", "-c", f"tests/resources/configs/manufactured-finetune_{model_name}_metrics_from_file.yaml"]
_ = build_lightning_cli(command_list)
gc.collect()