-
Notifications
You must be signed in to change notification settings - Fork 56
/
KLLoss.py
28 lines (24 loc) · 1.09 KB
/
KLLoss.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
# Code for "ActionCLIP: ActionCLIP: A New Paradigm for Action Recognition"
# arXiv:
# Mengmeng Wang, Jiazheng Xing, Yong Liu
import torch.nn.functional as F
import torch.nn as nn
class KLLoss(nn.Module):
"""Loss that uses a 'hinge' on the lower bound.
This means that for samples with a label value smaller than the threshold, the loss is zero if the prediction is
also smaller than that threshold.
args:
error_matric: What base loss to use (MSE by default).
threshold: Threshold to use for the hinge.
clip: Clip the loss if it is above this value.
"""
def __init__(self, error_metric=nn.KLDivLoss(size_average=True, reduce=True)):
super().__init__()
print('=========using KL Loss=and has temperature and * bz==========')
self.error_metric = error_metric
def forward(self, prediction, label):
batch_size = prediction.shape[0]
probs1 = F.log_softmax(prediction, 1)
probs2 = F.softmax(label * 10, 1)
loss = self.error_metric(probs1, probs2) * batch_size
return loss