206 lines
7.7 KiB
Python
Executable File
206 lines
7.7 KiB
Python
Executable File
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from criterions.loss_builder import CRITERIONS
|
|
|
|
|
|
@CRITERIONS.register('NONUNIFORM')
|
|
class OadLoss(nn.Module):
|
|
|
|
def __init__(self, cfg, reduction='mean'):
|
|
super(OadLoss, self).__init__()
|
|
self.reduction = reduction
|
|
self.num_classes = cfg['num_classes']
|
|
self.loss_offset = cfg['loss_offset'] if 'loss_offset' in cfg else 0
|
|
assert self.loss_offset==0, f'loss_offset is not 0: {self.loss_offset}'
|
|
self.loss = self.end_loss
|
|
|
|
def end_loss(self, out_dict, target):
|
|
# logits: (B, seq, K) target: (B, seq, K)
|
|
assert self.loss_offset < out_dict['logits'].shape[1], f'loss_offset is longer than the sequence loss_offset: {self.loss_offset}, logits.shape: {out_dict["logits"].shape}'
|
|
logits = out_dict['logits']
|
|
logits = logits[:,-1 - self.loss_offset,:].contiguous()
|
|
target = target[:,-1 - self.loss_offset,:].contiguous()
|
|
ce_loss = self.mlce_loss(logits, target)
|
|
return ce_loss
|
|
|
|
def mlce_loss(self, logits, target):
|
|
'''
|
|
multi label cross entropy loss.
|
|
logits: (B, K) target: (B, K)
|
|
'''
|
|
logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device)
|
|
output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B
|
|
if self.reduction == 'mean':
|
|
loss = torch.mean(output)
|
|
elif self.reduction == 'sum':
|
|
loss = torch.sum(output)
|
|
return loss
|
|
|
|
def forward(self, out_dict, target):
|
|
return self.loss(out_dict, target)
|
|
|
|
|
|
@CRITERIONS.register('UNIFORM')
|
|
class OadUniformLoss(nn.Module):
|
|
|
|
def __init__(self, cfg, reduction='mean'):
|
|
super(OadUniformLoss, self).__init__()
|
|
self.reduction = reduction
|
|
self.num_classes = cfg['num_classes']
|
|
# self.loss_offset = cfg['loss_offset'] if 'loss_offset' in cfg else 0
|
|
self.loss = self.end_loss
|
|
|
|
def end_loss(self, out_dict, target):
|
|
# logits: (B, seq, K) target: (B, seq, K)
|
|
logits = out_dict['logits']
|
|
# logits = logits[:,-1 - self.loss_offset,:].contiguous()
|
|
# target = target[:,-1 - self.loss_offset,:].contiguous()
|
|
ce_loss = self.mlce_loss(logits, target)
|
|
return ce_loss
|
|
|
|
def mlce_loss(self, logits, target):
|
|
'''
|
|
multi label cross entropy loss.
|
|
logits: (B, K) target: (B, K)
|
|
'''
|
|
logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device)
|
|
output = torch.sum(-target * logsoftmax(logits), dim=1) # B
|
|
# output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B
|
|
if self.reduction == 'mean':
|
|
loss = torch.mean(output)
|
|
elif self.reduction == 'sum':
|
|
loss = torch.sum(output)
|
|
return loss
|
|
|
|
def forward(self, out_dict, target):
|
|
return self.loss(out_dict, target)
|
|
|
|
@CRITERIONS.register('OADTR')
|
|
class OadTRLoss(nn.Module):
|
|
|
|
def __init__(self, cfg, reduction='mean'):
|
|
super(OadTRLoss, self).__init__()
|
|
self.reduction = reduction
|
|
self.num_classes = cfg['num_classes']
|
|
# self.loss_offset = cfg['loss_offset'] if 'loss_offset' in cfg else 0
|
|
self.loss = self.end_loss
|
|
|
|
def end_loss(self, out_dict, target):
|
|
# logits: (B, seq, K) target: (B, seq, K)
|
|
assert self.loss_offset < out_dict['logits'].shape[1], f'loss_offset is longer than the sequence loss_offset: {self.loss_offset}, logits.shape: {out_dict["logits"].shape}'
|
|
logits = out_dict['logits']
|
|
# logits = logits[:,-1 - self.loss_offset,:].contiguous()
|
|
# target = target[:,-1 - self.loss_offset,:].contiguous()
|
|
ce_loss = self.mlce_loss(logits, target)
|
|
return ce_loss
|
|
|
|
def mlce_loss(self, logits, target):
|
|
'''
|
|
multi label cross entropy loss.
|
|
logits: (B, K) target: (B, K)
|
|
'''
|
|
logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device)
|
|
output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B
|
|
if self.reduction == 'mean':
|
|
loss = torch.mean(output)
|
|
elif self.reduction == 'sum':
|
|
loss = torch.sum(output)
|
|
return loss
|
|
|
|
def forward(self, out_dict, enc_target, dec_target):
|
|
return self.loss(out_dict, enc_target, dec_target)
|
|
|
|
@CRITERIONS.register('ANTICIPATION')
|
|
class OadAntLoss(nn.Module):
|
|
|
|
# reduction was sum, not mean, what if we change to sum?
|
|
def __init__(self, cfg, reduction='mean'):
|
|
super(OadAntLoss, self).__init__()
|
|
self.reduction = reduction
|
|
self.loss = self.anticipation_loss
|
|
self.num_classes = cfg['num_classes']
|
|
|
|
def anticipation_loss(self, out_dict, target, ant_target):
|
|
# anticipation_logits: (B, seq, ANTICIPATION_LENGTH, num_class)
|
|
# ant_target: (B, ANTICIPATION_LENGTH, num_class)
|
|
anticipation_logits = out_dict['anticipation_logits']
|
|
# [B, window_size, anticipation_length, num_class]
|
|
pred_anticipation_logits = anticipation_logits[:,-1,:,:].contiguous().view(-1, self.num_classes)
|
|
|
|
anticipation_logit_targets = ant_target.view(-1, self.num_classes)
|
|
ant_loss = self.mlce_loss(pred_anticipation_logits, anticipation_logit_targets)
|
|
|
|
return ant_loss
|
|
|
|
def ce_loss(self, out_dict, target):
|
|
# logits: (B, seq, K) target: (B, seq, K)
|
|
logits = out_dict['logits']
|
|
logits = logits[:,-1,:].contiguous()
|
|
target = target[:,-1,:].contiguous()
|
|
ce_loss = self.mlce_loss(logits, target)
|
|
return ce_loss
|
|
|
|
def mlce_loss(self, logits, target):
|
|
'''
|
|
multi label cross entropy loss.
|
|
logits: (B, K) target: (B, K)
|
|
'''
|
|
logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device)
|
|
output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B
|
|
if self.reduction == 'mean':
|
|
loss = torch.mean(output)
|
|
elif self.reduction == 'sum':
|
|
loss = torch.sum(output)
|
|
|
|
return loss
|
|
|
|
def forward(self, out_dict, target, ant_target):
|
|
return self.loss(out_dict, target, ant_target)
|
|
|
|
|
|
|
|
@CRITERIONS.register('UNIFORM_ANTICIPATION')
|
|
class OadUniformAntLoss(nn.Module):
|
|
|
|
# reduction was sum, not mean, what if we change to sum?
|
|
def __init__(self, cfg, reduction='mean'):
|
|
super(OadUniformAntLoss, self).__init__()
|
|
self.reduction = reduction
|
|
self.loss = self.anticipation_loss
|
|
self.num_classes = cfg['num_classes']
|
|
|
|
def anticipation_loss(self, out_dict, target, ant_target):
|
|
anticipation_logits = out_dict['anticipation_logits']
|
|
pred_anticipation_logits = anticipation_logits[:,-1,:,:].contiguous().view(-1, self.num_classes)
|
|
anticipation_logit_targets = ant_target.view(-1, self.num_classes)
|
|
ant_loss = self.mlce_loss(pred_anticipation_logits, anticipation_logit_targets)
|
|
return ant_loss
|
|
|
|
def ce_loss(self, out_dict, target):
|
|
# logits: (B, seq, K) target: (B, seq, K)
|
|
logits = out_dict['logits']
|
|
logits = logits[:,-1,:].contiguous()
|
|
target = target[:,-1,:].contiguous()
|
|
ce_loss = self.mlce_loss(logits, target)
|
|
return ce_loss
|
|
|
|
def mlce_loss(self, logits, target):
|
|
'''
|
|
multi label cross entropy loss.
|
|
logits: (B, K) target: (B, K)
|
|
'''
|
|
logsoftmax = nn.LogSoftmax(dim=-1).to(logits.device)
|
|
output = torch.sum(-target * logsoftmax(logits), dim=1) # B
|
|
|
|
# output = torch.sum(-F.normalize(target) * logsoftmax(logits), dim=1) # B
|
|
if self.reduction == 'mean':
|
|
loss = torch.mean(output)
|
|
elif self.reduction == 'sum':
|
|
loss = torch.sum(output)
|
|
|
|
return loss
|
|
|
|
def forward(self, out_dict, target, ant_target):
|
|
return self.loss(out_dict, target, ant_target)
|