-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathvicReg.py
More file actions
108 lines (81 loc) · 3.98 KB
/
vicReg.py
File metadata and controls
108 lines (81 loc) · 3.98 KB
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
# Copyright 2021 solo-learn development team.
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
# Software, and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all copies
# or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
# FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.
import torch
import torch.nn.functional as F
def invariance_loss(z1: torch.Tensor, z2: torch.Tensor) -> torch.Tensor:
"""Computes mse loss given batch of projected features z1 from view 1 and
projected features z2 from view 2.
Args:
z1 (torch.Tensor): NxD Tensor containing projected features from view 1.
z2 (torch.Tensor): NxD Tensor containing projected features from view 2.
Returns:
torch.Tensor: invariance loss (mean squared error).
"""
return F.mse_loss(z1, z2)
def variance_loss(z1: torch.Tensor, z2: torch.Tensor) -> torch.Tensor:
"""Computes variance loss given batch of projected features z1 from view 1 and
projected features z2 from view 2.
Args:
z1 (torch.Tensor): NxD Tensor containing projected features from view 1.
z2 (torch.Tensor): NxD Tensor containing projected features from view 2.
Returns:
torch.Tensor: variance regularization loss.
"""
eps = 1e-4
std_z1 = torch.sqrt(z1.var(dim=0) + eps)
std_z2 = torch.sqrt(z2.var(dim=0) + eps)
std_loss = torch.mean(F.relu(.8 - std_z1)) + torch.mean(F.relu(.8 - std_z2))
return std_loss
def covariance_loss(z1: torch.Tensor, z2: torch.Tensor) -> torch.Tensor:
"""Computes covariance loss given batch of projected features z1 from view 1 and
projected features z2 from view 2.
Args:
z1 (torch.Tensor): NxD Tensor containing projected features from view 1.
z2 (torch.Tensor): NxD Tensor containing projected features from view 2.
Returns:
torch.Tensor: covariance regularization loss.
"""
N, D = z1.size()
z1 = z1 - z1.mean(dim=0)
z2 = z2 - z2.mean(dim=0)
cov_z1 = (z1.T @ z1) / (N - 1)
cov_z2 = (z2.T @ z2) / (N - 1)
diag = torch.eye(D, device=z1.device)
cov_loss = cov_z1[~diag.bool()].pow_(2).sum() / D + cov_z2[~diag.bool()].pow_(2).sum() / D
return cov_loss
def vicreg_loss_func(
z1: torch.Tensor,
z2: torch.Tensor,
sim_loss_weight: float = 25.0,
var_loss_weight: float = 25.0,
cov_loss_weight: float = 1.0,
) -> torch.Tensor:
"""Computes VICReg's loss given batch of projected features z1 from view 1 and
projected features z2 from view 2.
Args:
z1 (torch.Tensor): NxD Tensor containing projected features from view 1.
z2 (torch.Tensor): NxD Tensor containing projected features from view 2.
sim_loss_weight (float): invariance loss weight.
var_loss_weight (float): variance loss weight.
cov_loss_weight (float): covariance loss weight.
Returns:
torch.Tensor: VICReg loss.
"""
sim_loss = invariance_loss(z1, z2)
var_loss = variance_loss(z1, z2)
cov_loss = covariance_loss(z1, z2)
loss = sim_loss_weight * sim_loss + var_loss_weight * var_loss + cov_loss_weight * cov_loss
return loss