KAEL_WSDM23/model.py at main · DEEP-PolyU/KAEL_WSDM23 · 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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
from threading import Thread
from tqdm import tqdm, trange
import utils as dp
import logging
import param


class MAB():

    def __init__(self, opportunities, dictionary, detected, θ, ita) -> None:
        self.oppoetunities = opportunities
        self.dictionary = dictionary
        self.detected = detected
        self.θ = θ
        self.ita = ita

    def initialize(self, overlaps, para_lambda4):
        """
        Initialize the parameters through overlaps

        Returns:
            [θ]: [Initialized θ, 300 dimensional tensor]
            [ita]: [Initialized ita, 300 dimensional value]
        """
        print('**************Start Initialization**************')
        total_rewards = 0  # Initialize total rewards
        rewards = [[]] * param.k  # Initialize rewards
        Y = [[]] * param.k
        for ini in tqdm(range(param.k)):
            for i in range(len(overlaps)):
                #print('------ Iteration ------',i)
                reward = dp.rewarding(
                    self.dictionary, overlaps[i]
                )  # Check label by rewarding, 1 means anomaly, 0 means nominal
                rewards[ini].append(
                    reward
                )  # Update y, which indicates rewards here for current cluster
                Y[ini].append(overlaps[i])  # Update Y for current cluster
                self.θ[ini] = dp.θ_update_torch(
                    Y[ini], rewards[ini],
                    para_lambda4[ini])  # Update θ for current cluster
                self.ita[ini] = dp.ita_update_torch(
                        Y[ini], overlaps[i],
                    para_lambda4[ini])  # Update ita for current cluster
                total_rewards += reward
                #**************************************************#
        print('*************************')
        print('❀❀❀❀❀❀❀❀❀❀ Initialized successfully ❀❀❀❀❀❀❀❀❀❀')
        return self.θ, self.ita, rewards, Y

    def ranking(self, arm, triples, detected, besttriples, expectation):
        for i in range(len(triples)):
            if triples[i] not in detected:
                # Calculate the expectation
                expectationtemp = dp.calExpectation_torch(
                    triples[i], self.θ[arm], self.ita[arm])
                # Store the max expectation
                if expectationtemp >= expectation[arm]:
                    besttriples[arm] = triples[
                        i]  # Record the triple with highest expectation for current cluster
                    expectation[
                        arm] = expectationtemp  # Record the best expectation for current cluster

    def train(self, All_triples):
        global ablationcount
        global ablationrecord
        global ablationrewards
        global armchoice
        total_rewards = 0  # Initialize total rewards
        Y = [[]] * param.k
        rewards = [[] for _ in range(param.k)]  # Initialize rewards
        expectation = [0] * param.k  # Initialize Expectation
        besttriples = [None] * param.k  # Initialize the best triples for each arm
        threads = []  # Initialize the threads list

        for i in range(self.opportunities):
            print('------Oracle Iteration ------', i)
            # Create and start threads for each arm
            for arm in range(param.k):
                thread = Thread(target=self.ranking,
                                args=(arm, All_triples[arm], self.detected,
                                      besttriples, expectation))
                threads.append(thread)
                thread.start()

            for thread in threads:
                thread.join()

            #**************************************************#
            bestarm = expectation.index(
                max(expectation))  # Find the best among 3
            reward = dp.rewarding(
                self.dictionary, besttriples[bestarm]
            )  # Check label by rewarding, 1 means anomaly, -0.01 means nominal
            rewards[bestarm].append(
                reward
            )  # Update y, which indicates rewards here for current cluster
            Y[bestarm].append(
                besttriples[bestarm])  # Update Y for current cluster
            self.θ[bestarm] = dp.θ_update_torch(
                Y[bestarm], rewards[bestarm],
                param.para_lambda4[bestarm])  # Update θ for current cluster
            self.ita[bestarm] = dp.ita_update_torch(
                Y[bestarm], besttriples[bestarm],
                param.para_lambda4[bestarm])  # Update ita for current cluster
            total_rewards += reward
            self.detected.append(besttriples[bestarm])

            print('***** Current TNR For arm: *****', bestarm)
            print(total_rewards / (i + 1))

        TNR = dp.calTNR(total_rewards, self.opportunities)
        return total_rewards, TNR, self.θ, self.ita, rewards, Y

    def application(self, iteration, All_triples):
        global ablationcount
        global ablationrecord
        global ablationrewards
        global armchoice
        total_rewards = 0
        besttriples = [[]] * param.k  # Initialize best choices of each cluster in in current iteration
        expectation = [0] * param.k  # Initialize Expectation
        for i in trange(iteration):
            print('------ Iteration ------', i)
            expectation = [0] * param.k
            threads = []
            for j in range(param.k):
                thread = Thread(target=self.ranking, args=(j, All_triples[j]))
                threads.append(thread)
                thread.start()
            for thread in threads:
                thread.join()
            bestarm = expectation.index(max(expectation))
            reward = dp.rewarding(self.dictionary, besttriples[bestarm])
            total_rewards += reward
            self.detected.append(besttriples[bestarm])
            print('***** Current TNR For arm: *****', bestarm)
            print(total_rewards / (i + 1))
        TNR = dp.calTNR(total_rewards, iteration)
        return total_rewards, TNR, self.θ, self.ita