-
Notifications
You must be signed in to change notification settings - Fork 9
/
regularizers.py
142 lines (108 loc) · 3.66 KB
/
regularizers.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
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
141
142
from abc import ABC, abstractmethod
from typing import Tuple
import torch
from torch import nn
class Regularizer(nn.Module, ABC):
@abstractmethod
def forward(self, factors: Tuple[torch.Tensor]):
pass
class Fro(Regularizer):
def __init__(self, weight: float):
super(Fro, self).__init__()
self.weight = weight
def forward(self, factors):
norm = 0
for factor in factors:
for f in factor:
norm += self.weight * torch.sum(
torch.norm(f, 2) ** 2
)
return norm / factors[0][0].shape[0]
class N3(Regularizer):
def __init__(self, weight: float):
super(N3, self).__init__()
self.weight = weight
def forward(self, factors):
norm = 0
for factor in factors:
for f in factor:
norm += self.weight * torch.sum(
torch.abs(f) ** 3
) / f.shape[0]
return norm
class L2(Regularizer):
def __init__(self, weight: float):
super(L2, self).__init__()
self.weight = weight
def forward(self, factors):
norm = 0
for factor in factors:
for f in factor:
norm += self.weight * torch.sum(
torch.abs(f) ** 2
)
return norm / factors[0][0].shape[0]
class L1(Regularizer):
def __init__(self, weight: float):
super(L1, self).__init__()
self.weight = weight
def forward(self, factors):
norm = 0
for factor in factors:
for f in factor:
norm += self.weight * torch.sum(
torch.abs(f)**1
)
return norm / factors[0][0].shape[0]
class NA(Regularizer):
def __init__(self, weight: float):
super(NA, self).__init__()
self.weight = weight
def forward(self, factors):
return torch.Tensor([0.0]).cuda()
class DURA(Regularizer):
def __init__(self, weight: float):
super(DURA, self).__init__()
self.weight = weight
def forward(self, factors):
norm = 0
for factor in factors:
h, r, t = factor
norm += torch.sum(t**2 + h**2)
norm += torch.sum(h**2 * r**2 + t**2 * r**2)
return self.weight * norm / h.shape[0]
class DURA_RESCAL(Regularizer):
def __init__(self, weight: float):
super(DURA_RESCAL, self).__init__()
self.weight = weight
def forward(self, factors):
norm = 0
for factor in factors:
h, r, t = factor
norm += torch.sum(h ** 2 + t ** 2)
norm += torch.sum(
torch.bmm(r.transpose(1, 2), h.unsqueeze(-1)) ** 2 + torch.bmm(r, t.unsqueeze(-1)) ** 2)
return self.weight * norm / h.shape[0]
class DURA_RESCAL_W(Regularizer):
def __init__(self, weight: float):
super(DURA_RESCAL_W, self).__init__()
self.weight = weight
def forward(self, factors):
norm = 0
for factor in factors:
h, r, t = factor
norm += 2.0 * torch.sum(h ** 2 + t ** 2)
norm += 0.5 * torch.sum(
torch.bmm(r.transpose(1, 2), h.unsqueeze(-1)) ** 2 + torch.bmm(r, t.unsqueeze(-1)) ** 2)
return self.weight * norm / h.shape[0]
class DURA_W(Regularizer):
def __init__(self, weight: float):
super(DURA_W, self).__init__()
self.weight = weight
def forward(self, factors):
norm = 0
for factor in factors:
h, r, t = factor
norm += 0.5 * torch.sum(t**2 + h**2)
norm += 1.5 * torch.sum(h**2 * r**2 + t**2 * r**2)
return self.weight * norm / h.shape[0]