dice-embeddings/tests/test_adaptive_swa.py at develop · dice-group/dice-embeddings · 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
109
110
111
from dicee.executer import Execute
import pytest
from dicee.config import Namespace


class TestASWA:

    @pytest.mark.filterwarnings('ignore::UserWarning')
    def test_k_vs_all_lowest(self):
        args = Namespace()
        args.model = 'Keci'
        args.p = 0
        args.q = 1
        args.scoring_technique = "KvsAll"
        args.dataset_dir = "KGs/UMLS"
        args.trainer = "PL"
        args.num_epochs = 100
        args.lr = 0.1
        args.embedding_dim = 32
        args.batch_size = 1024
        args.adaptive_swa = True
        aswa_report = Execute(args).start()

        args = Namespace()
        args.model = 'Keci'
        args.p = 0
        args.q = 1
        args.scoring_technique = "KvsAll"
        args.dataset_dir = "KGs/UMLS"
        args.trainer = "PL"
        args.num_epochs = 100
        args.lr = 0.1
        args.embedding_dim = 32
        args.batch_size = 1024
        args.swa = True
        swa_report = Execute(args).start()

        assert aswa_report["Val"]["MRR"] > swa_report["Val"]["MRR"]
        assert aswa_report["Val"]["H@1"] > swa_report["Val"]["H@1"]
        assert aswa_report["Test"]["MRR"] > swa_report["Test"]["MRR"]
        assert aswa_report["Test"]["H@1"] > swa_report["Test"]["H@1"]

    """

    @pytest.mark.filterwarnings('ignore::UserWarning')
    def test_k_vs_all_low(self):
        args = Namespace()
        args.model = 'Keci'
        args.p = 0
        args.q = 1
        args.scoring_technique = "KvsAll"
        args.dataset_dir = "KGs/UMLS"
        args.trainer = "PL"
        args.num_epochs = 50
        args.lr = 0.1
        args.embedding_dim = 32
        args.batch_size = 1024
        args.adaptive_swa = True
        aswa_report = Execute(args).start()

        args = Namespace()
        args.model = 'Keci'
        args.p = 0
        args.q = 1
        args.scoring_technique = "KvsAll"
        args.dataset_dir = "KGs/UMLS"
        args.trainer = "PL"
        args.num_epochs = 50
        args.lr = 0.1
        args.embedding_dim = 32
        args.batch_size = 1024
        args.stochastic_weight_avg = True
        swa_report = Execute(args).start()

        assert aswa_report["Test"]["MRR"] > swa_report["Test"]["MRR"]
        assert aswa_report["Test"]["H@1"] > swa_report["Test"]["H@1"]

    @pytest.mark.filterwarnings('ignore::UserWarning')
    def test_k_vs_all_mid(self):
        args = Namespace()
        args.model = 'Keci'
        args.p = 0
        args.q = 1
        args.scoring_technique = "KvsAll"
        args.dataset_dir = "KGs/UMLS"
        args.trainer = "PL"
        args.num_epochs = 200
        args.lr = 0.1
        args.embedding_dim = 32
        args.batch_size = 1024
        args.adaptive_swa = True
        aswa_report = Execute(args).start()

        args = Namespace()
        args.model = 'Keci'
        args.p = 0
        args.q = 1
        args.scoring_technique = "KvsAll"
        args.dataset_dir = "KGs/UMLS"
        args.trainer = "PL"
        args.num_epochs = 200
        args.lr = 0.1
        args.embedding_dim = 32
        args.batch_size = 1024
        args.stochastic_weight_avg = True
        swa_report = Execute(args).start()

        assert aswa_report["Val"]["MRR"] > swa_report["Val"]["MRR"]
        assert aswa_report["Test"]["MRR"] > swa_report["Test"]["MRR"]
        assert 0.88 > aswa_report["Test"]["MRR"] > swa_report["Test"]["MRR"] > 0.75
    """